ha-core/homeassistant/components/waqi/sensor.py

"""Support for the World Air Quality Index service."""
from __future__ import annotations

import asyncio
from datetime import timedelta
import logging

import aiohttp
import voluptuous as vol
from waqiasync import WaqiClient

from homeassistant.components.sensor import (
    SensorDeviceClass,
    SensorEntity,
    SensorStateClass,
)
from homeassistant.const import (
    ATTR_ATTRIBUTION,
    ATTR_TEMPERATURE,
    ATTR_TIME,
    CONF_TOKEN,
)
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import PlatformNotReady
from homeassistant.helpers.aiohttp_client import async_get_clientsession
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.config_validation import PLATFORM_SCHEMA
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType

_LOGGER = logging.getLogger(__name__)

ATTR_DOMINENTPOL = "dominentpol"
ATTR_HUMIDITY = "humidity"
ATTR_NITROGEN_DIOXIDE = "nitrogen_dioxide"
ATTR_OZONE = "ozone"
ATTR_PM10 = "pm_10"
ATTR_PM2_5 = "pm_2_5"
ATTR_PRESSURE = "pressure"
ATTR_SULFUR_DIOXIDE = "sulfur_dioxide"

KEY_TO_ATTR = {
    "pm25": ATTR_PM2_5,
    "pm10": ATTR_PM10,
    "h": ATTR_HUMIDITY,
    "p": ATTR_PRESSURE,
    "t": ATTR_TEMPERATURE,
    "o3": ATTR_OZONE,
    "no2": ATTR_NITROGEN_DIOXIDE,
    "so2": ATTR_SULFUR_DIOXIDE,
}

ATTRIBUTION = "Data provided by the World Air Quality Index project"

ATTR_ICON = "mdi:cloud"
ATTR_UNIT = "AQI"

CONF_LOCATIONS = "locations"
CONF_STATIONS = "stations"

SCAN_INTERVAL = timedelta(minutes=5)

TIMEOUT = 10

PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
    {
        vol.Optional(CONF_STATIONS): cv.ensure_list,
        vol.Required(CONF_TOKEN): cv.string,
        vol.Required(CONF_LOCATIONS): cv.ensure_list,
    }
)


async def async_setup_platform(
    hass: HomeAssistant,
    config: ConfigType,
    async_add_entities: AddEntitiesCallback,
    discovery_info: DiscoveryInfoType | None = None,
) -> None:
    """Set up the requested World Air Quality Index locations."""

    token = config[CONF_TOKEN]
    station_filter = config.get(CONF_STATIONS)
    locations = config[CONF_LOCATIONS]

    client = WaqiClient(token, async_get_clientsession(hass), timeout=TIMEOUT)
    dev = []
    try:
        for location_name in locations:
            stations = await client.search(location_name)
            _LOGGER.debug("The following stations were returned: %s", stations)
            for station in stations:
                waqi_sensor = WaqiSensor(client, station)
                if not station_filter or {
                    waqi_sensor.uid,
                    waqi_sensor.url,
                    waqi_sensor.station_name,
                } & set(station_filter):
                    dev.append(waqi_sensor)
    except (
        aiohttp.client_exceptions.ClientConnectorError,
        asyncio.TimeoutError,
    ) as err:
        _LOGGER.exception("Failed to connect to WAQI servers")
        raise PlatformNotReady from err
    async_add_entities(dev, True)


class WaqiSensor(SensorEntity):
    """Implementation of a WAQI sensor."""

    _attr_icon = ATTR_ICON
    _attr_native_unit_of_measurement = ATTR_UNIT
    _attr_device_class = SensorDeviceClass.AQI
    _attr_state_class = SensorStateClass.MEASUREMENT

    def __init__(self, client, station):
        """Initialize the sensor."""
        self._client = client
        try:
            self.uid = station["uid"]
        except (KeyError, TypeError):
            self.uid = None

        try:
            self.url = station["station"]["url"]
        except (KeyError, TypeError):
            self.url = None

        try:
            self.station_name = station["station"]["name"]
        except (KeyError, TypeError):
            self.station_name = None

        self._data = None

    @property
    def name(self):
        """Return the name of the sensor."""
        if self.station_name:
            return f"WAQI {self.station_name}"
        return f"WAQI {self.url if self.url else self.uid}"

    @property
    def native_value(self):
        """Return the state of the device."""
        if self._data is not None:
            return self._data.get("aqi")
        return None

    @property
    def available(self):
        """Return sensor availability."""
        return self._data is not None

    @property
    def unique_id(self):
        """Return unique ID."""
        return self.uid

    @property
    def extra_state_attributes(self):
        """Return the state attributes of the last update."""
        attrs = {}

        if self._data is not None:
            try:
                attrs[ATTR_ATTRIBUTION] = " and ".join(
                    [ATTRIBUTION]
                    + [v["name"] for v in self._data.get("attributions", [])]
                )

                attrs[ATTR_TIME] = self._data["time"]["s"]
                attrs[ATTR_DOMINENTPOL] = self._data.get("dominentpol")

                iaqi = self._data["iaqi"]
                for key in iaqi:
                    if key in KEY_TO_ATTR:
                        attrs[KEY_TO_ATTR[key]] = iaqi[key]["v"]
                    else:
                        attrs[key] = iaqi[key]["v"]
                return attrs
            except (IndexError, KeyError):
                return {ATTR_ATTRIBUTION: ATTRIBUTION}

    async def async_update(self):
        """Get the latest data and updates the states."""
        if self.uid:
            result = await self._client.get_station_by_number(self.uid)
        elif self.url:
            result = await self._client.get_station_by_name(self.url)
        else:
            result = None
        self._data = result