Move WAQI state attributes to separate sensors (home-assistant#101217)

* Migrate WAQI to has entity name

* Split WAQI extra state attributes into separate sensors

* Split WAQI extra state attributes into separate sensors

* Fix test

* Support new aiowaqi

* Bump aiowaqi to 2.1.0

* Add nephelometry as possible value

* Fix test

homeassistant/components/waqi/sensor.py

@@ -1,24 +1,36 @@
 """Support for the World Air Quality Index service."""
 from __future__ import annotations
 
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
 import logging
+from typing import Any
 
-from aiowaqi import WAQIAuthenticationError, WAQIClient, WAQIConnectionError
+from aiowaqi import (
+    WAQIAirQuality,
+    WAQIAuthenticationError,
+    WAQIClient,
+    WAQIConnectionError,
+)
+from aiowaqi.models import Pollutant
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
+    SensorEntityDescription,
     SensorStateClass,
 )
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import (
-    ATTR_ATTRIBUTION,
     ATTR_TEMPERATURE,
     ATTR_TIME,
     CONF_API_KEY,
     CONF_NAME,
     CONF_TOKEN,
+    PERCENTAGE,
+    UnitOfPressure,
+    UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import PlatformNotReady
@@ -27,7 +39,7 @@
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
-from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import CONF_STATION_NUMBER, DOMAIN, ISSUE_PLACEHOLDER
@@ -141,67 +153,167 @@ async def async_setup_platform(
         )
 
 
+@dataclass
+class WAQIMixin:
+    """Mixin for required keys."""
+
+    available_fn: Callable[[WAQIAirQuality], bool]
+    value_fn: Callable[[WAQIAirQuality], StateType]
+
+
+@dataclass
+class WAQISensorEntityDescription(SensorEntityDescription, WAQIMixin):
+    """Describes WAQI sensor entity."""
+
+
+SENSORS: list[WAQISensorEntityDescription] = [
+    WAQISensorEntityDescription(
+        key="air_quality",
+        device_class=SensorDeviceClass.AQI,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.air_quality_index,
+        available_fn=lambda _: True,
+    ),
+    WAQISensorEntityDescription(
+        key="humidity",
+        device_class=SensorDeviceClass.HUMIDITY,
+        native_unit_of_measurement=PERCENTAGE,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.humidity,
+        available_fn=lambda aq: aq.extended_air_quality.humidity is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="pressure",
+        device_class=SensorDeviceClass.PRESSURE,
+        native_unit_of_measurement=UnitOfPressure.HPA,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.pressure,
+        available_fn=lambda aq: aq.extended_air_quality.pressure is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="temperature",
+        device_class=SensorDeviceClass.TEMPERATURE,
+        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.temperature,
+        available_fn=lambda aq: aq.extended_air_quality.temperature is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="carbon_monoxide",
+        translation_key="carbon_monoxide",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.carbon_monoxide,
+        available_fn=lambda aq: aq.extended_air_quality.carbon_monoxide is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="nitrogen_dioxide",
+        translation_key="nitrogen_dioxide",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.nitrogen_dioxide,
+        available_fn=lambda aq: aq.extended_air_quality.nitrogen_dioxide is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="ozone",
+        translation_key="ozone",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.ozone,
+        available_fn=lambda aq: aq.extended_air_quality.ozone is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="sulphur_dioxide",
+        translation_key="sulphur_dioxide",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.sulfur_dioxide,
+        available_fn=lambda aq: aq.extended_air_quality.sulfur_dioxide is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="pm10",
+        translation_key="pm10",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.pm10,
+        available_fn=lambda aq: aq.extended_air_quality.pm10 is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="pm25",
+        translation_key="pm25",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.pm25,
+        available_fn=lambda aq: aq.extended_air_quality.pm25 is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="dominant_pollutant",
+        translation_key="dominant_pollutant",
+        device_class=SensorDeviceClass.ENUM,
+        options=[pollutant.value for pollutant in Pollutant],
+        value_fn=lambda aq: aq.dominant_pollutant,
+        available_fn=lambda _: True,
+    ),
+]
+
+
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the WAQI sensor."""
     coordinator: WAQIDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([WaqiSensor(coordinator)])
+    async_add_entities(
+        [
+            WaqiSensor(coordinator, sensor)
+            for sensor in SENSORS
+            if sensor.available_fn(coordinator.data)
+        ]
+    )
 
 
 class WaqiSensor(CoordinatorEntity[WAQIDataUpdateCoordinator], SensorEntity):
     """Implementation of a WAQI sensor."""
 
-    _attr_icon = ATTR_ICON
-    _attr_device_class = SensorDeviceClass.AQI
-    _attr_state_class = SensorStateClass.MEASUREMENT
     _attr_has_entity_name = True
-    _attr_name = None
+    entity_description: WAQISensorEntityDescription
 
-    def __init__(self, coordinator: WAQIDataUpdateCoordinator) -> None:
+    def __init__(
+        self,
+        coordinator: WAQIDataUpdateCoordinator,
+        entity_description: WAQISensorEntityDescription,
+    ) -> None:
         """Initialize the sensor."""
         super().__init__(coordinator)
-        self._attr_unique_id = f"{coordinator.data.station_id}_air_quality"
+        self.entity_description = entity_description
+        self._attr_unique_id = f"{coordinator.data.station_id}_{entity_description.key}"
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, str(coordinator.data.station_id))},
             name=coordinator.data.city.name,
             entry_type=DeviceEntryType.SERVICE,
         )
+        self._attr_attribution = " and ".join(
+            attribution.name for attribution in coordinator.data.attributions
+        )
 
     @property
-    def native_value(self) -> int | None:
+    def native_value(self) -> StateType:
         """Return the state of the device."""
-        return self.coordinator.data.air_quality_index
+        return self.entity_description.value_fn(self.coordinator.data)
 
     @property
-    def extra_state_attributes(self):
-        """Return the state attributes of the last update."""
-        attrs = {}
-        try:
-            attrs[ATTR_ATTRIBUTION] = " and ".join(
-                [ATTRIBUTION]
-                + [
-                    attribution.name
-                    for attribution in self.coordinator.data.attributions
-                ]
-            )
-
-            attrs[ATTR_TIME] = self.coordinator.data.measured_at
-            attrs[ATTR_DOMINENTPOL] = self.coordinator.data.dominant_pollutant
-
-            iaqi = self.coordinator.data.extended_air_quality
-
-            attribute = {
-                ATTR_PM2_5: iaqi.pm25,
-                ATTR_PM10: iaqi.pm10,
-                ATTR_HUMIDITY: iaqi.humidity,
-                ATTR_PRESSURE: iaqi.pressure,
-                ATTR_TEMPERATURE: iaqi.temperature,
-                ATTR_OZONE: iaqi.ozone,
-                ATTR_NITROGEN_DIOXIDE: iaqi.nitrogen_dioxide,
-                ATTR_SULFUR_DIOXIDE: iaqi.sulfur_dioxide,
-            }
-            res_attributes = {k: v for k, v in attribute.items() if v is not None}
-            return {**attrs, **res_attributes}
-        except (IndexError, KeyError):
-            return {ATTR_ATTRIBUTION: ATTRIBUTION}
+    def extra_state_attributes(self) -> Mapping[str, Any] | None:
+        """Return old state attributes if the entity is AQI entity."""
+        if self.entity_description.key != "air_quality":
+            return None
+        attrs: dict[str, Any] = {}
+        attrs[ATTR_TIME] = self.coordinator.data.measured_at
+        attrs[ATTR_DOMINENTPOL] = self.coordinator.data.dominant_pollutant
+
+        iaqi = self.coordinator.data.extended_air_quality
+
+        attribute = {
+            ATTR_PM2_5: iaqi.pm25,
+            ATTR_PM10: iaqi.pm10,
+            ATTR_HUMIDITY: iaqi.humidity,
+            ATTR_PRESSURE: iaqi.pressure,
+            ATTR_TEMPERATURE: iaqi.temperature,
+            ATTR_OZONE: iaqi.ozone,
+            ATTR_NITROGEN_DIOXIDE: iaqi.nitrogen_dioxide,
+            ATTR_SULFUR_DIOXIDE: iaqi.sulfur_dioxide,
+        }
+        res_attributes = {k: v for k, v in attribute.items() if v is not None}
+        return {**attrs, **res_attributes}

homeassistant/components/waqi/strings.json

@@ -53,5 +53,39 @@
       "title": "The WAQI YAML configuration import failed",
       "description": "Configuring World Air Quality Index using YAML is being removed but there weren't any stations imported because they couldn't be found.\n\nEnsure the imported configuration is correct and remove the World Air Quality Index YAML configuration from your configuration.yaml file and continue to [set up the integration]({url}) manually."
     }
+  },
+  "entity": {
+    "sensor": {
+      "carbon_monoxide": {
+        "name": "[%key:component::sensor::entity_component::carbon_monoxide::name%]"
+      },
+      "nitrogen_dioxide": {
+        "name": "[%key:component::sensor::entity_component::nitrogen_dioxide::name%]"
+      },
+      "ozone": {
+        "name": "[%key:component::sensor::entity_component::ozone::name%]"
+      },
+      "sulphur_dioxide": {
+        "name": "[%key:component::sensor::entity_component::sulphur_dioxide::name%]"
+      },
+      "pm10": {
+        "name": "[%key:component::sensor::entity_component::pm10::name%]"
+      },
+      "pm25": {
+        "name": "[%key:component::sensor::entity_component::pm25::name%]"
+      },
+      "dominant_pollutant": {
+        "name": "Dominant pollutant",
+        "state": {
+          "co": "[%key:component::sensor::entity_component::carbon_monoxide::name%]",
+          "neph": "Nephelometry",
+          "no2": "[%key:component::sensor::entity_component::nitrogen_dioxide::name%]",
+          "o3": "[%key:component::sensor::entity_component::ozone::name%]",
+          "so2": "[%key:component::sensor::entity_component::sulphur_dioxide::name%]",
+          "pm10": "[%key:component::sensor::entity_component::pm10::name%]",
+          "pm25": "[%key:component::sensor::entity_component::pm25::name%]"
+        }
+      }
+    }
   }
 }

tests/components/waqi/fixtures/air_quality_sensor.json

@@ -23,9 +23,15 @@
     "h": {
       "v": 80
     },
+    "co": {
+      "v": 2.3
+    },
     "no2": {
       "v": 2.3
     },
+    "so2": {
+      "v": 2.3
+    },
     "o3": {
       "v": 29.4
     },