Use EntityDescription - trafikverket_weatherstation (#54430)

homeassistant/components/trafikverket_weatherstation/sensor.py

@@ -1,6 +1,8 @@
 """Weather information for air and road temperature (by Trafikverket)."""
+from __future__ import annotations
 
 import asyncio
+from dataclasses import dataclass
 from datetime import timedelta
 import logging
 
@@ -8,7 +10,11 @@ import aiohttp
 from pytrafikverket.trafikverket_weather import TrafikverketWeather
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.const import (
     ATTR_ATTRIBUTION,
     CONF_API_KEY,
@@ -38,85 +44,102 @@ MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=10)
 
 SCAN_INTERVAL = timedelta(seconds=300)
 
-SENSOR_TYPES = {
+
-    "air_temp": [
+@dataclass
-        "Air temperature",
+class TrafikverketRequiredKeysMixin:
-        TEMP_CELSIUS,
+    """Mixin for required keys."""
-        "air_temp",
+
-        "mdi:thermometer",
+    api_key: str
-        DEVICE_CLASS_TEMPERATURE,
+
-    ],
+
-    "road_temp": [
+@dataclass
-        "Road temperature",
+class TrafikverketSensorEntityDescription(
-        TEMP_CELSIUS,
+    SensorEntityDescription, TrafikverketRequiredKeysMixin
-        "road_temp",
+):
-        "mdi:thermometer",
+    """Describes Trafikverket sensor entity."""
-        DEVICE_CLASS_TEMPERATURE,
+
-    ],
+
-    "precipitation": [
+SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = (
-        "Precipitation type",
+    TrafikverketSensorEntityDescription(
-        None,
+        key="air_temp",
-        "precipitationtype",
+        api_key="air_temp",
-        "mdi:weather-snowy-rainy",
+        name="Air temperature",
-        None,
+        native_unit_of_measurement=TEMP_CELSIUS,
-    ],
+        icon="mdi:thermometer",
-    "wind_direction": [
+        device_class=DEVICE_CLASS_TEMPERATURE,
-        "Wind direction",
+    ),
-        DEGREE,
+    TrafikverketSensorEntityDescription(
-        "winddirection",
+        key="road_temp",
-        "mdi:flag-triangle",
+        api_key="road_temp",
-        None,
+        name="Road temperature",
-    ],
+        native_unit_of_measurement=TEMP_CELSIUS,
-    "wind_direction_text": [
+        icon="mdi:thermometer",
-        "Wind direction text",
+        device_class=DEVICE_CLASS_TEMPERATURE,
-        None,
+    ),
-        "winddirectiontext",
+    TrafikverketSensorEntityDescription(
-        "mdi:flag-triangle",
+        key="precipitation",
-        None,
+        api_key="precipitationtype",
-    ],
+        name="Precipitation type",
-    "wind_speed": [
+        icon="mdi:weather-snowy-rainy",
-        "Wind speed",
+    ),
-        SPEED_METERS_PER_SECOND,
+    TrafikverketSensorEntityDescription(
-        "windforce",
+        key="wind_direction",
-        "mdi:weather-windy",
+        api_key="winddirection",
-        None,
+        name="Wind direction",
-    ],
+        native_unit_of_measurement=DEGREE,
-    "wind_speed_max": [
+        icon="mdi:flag-triangle",
-        "Wind speed max",
+    ),
-        SPEED_METERS_PER_SECOND,
+    TrafikverketSensorEntityDescription(
-        "windforcemax",
+        key="wind_direction_text",
-        "mdi:weather-windy-variant",
+        api_key="winddirectiontext",
-        None,
+        name="Wind direction text",
-    ],
+        icon="mdi:flag-triangle",
-    "humidity": [
+    ),
-        "Humidity",
+    TrafikverketSensorEntityDescription(
-        PERCENTAGE,
+        key="wind_speed",
-        "humidity",
+        api_key="windforce",
-        "mdi:water-percent",
+        name="Wind speed",
-        DEVICE_CLASS_HUMIDITY,
+        native_unit_of_measurement=SPEED_METERS_PER_SECOND,
-    ],
+        icon="mdi:weather-windy",
-    "precipitation_amount": [
+    ),
-        "Precipitation amount",
+    TrafikverketSensorEntityDescription(
-        LENGTH_MILLIMETERS,
+        key="wind_speed_max",
-        "precipitation_amount",
+        api_key="windforcemax",
-        "mdi:cup-water",
+        name="Wind speed max",
-        None,
+        native_unit_of_measurement=SPEED_METERS_PER_SECOND,
-    ],
+        icon="mdi:weather-windy-variant",
-    "precipitation_amountname": [
+    ),
-        "Precipitation name",
+    TrafikverketSensorEntityDescription(
-        None,
+        key="humidity",
-        "precipitation_amountname",
+        api_key="humidity",
-        "mdi:weather-pouring",
+        name="Humidity",
-        None,
+        native_unit_of_measurement=PERCENTAGE,
-    ],
+        icon="mdi:water-percent",
-}
+        device_class=DEVICE_CLASS_HUMIDITY,
+    ),
+    TrafikverketSensorEntityDescription(
+        key="precipitation_amount",
+        api_key="precipitation_amount",
+        name="Precipitation amount",
+        native_unit_of_measurement=LENGTH_MILLIMETERS,
+        icon="mdi:cup-water",
+    ),
+    TrafikverketSensorEntityDescription(
+        key="precipitation_amountname",
+        api_key="precipitation_amountname",
+        name="Precipitation name",
+        icon="mdi:weather-pouring",
+    ),
+)
+
+SENSOR_KEYS = [desc.key for desc in SENSOR_TYPES]
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Required(CONF_NAME): cv.string,
         vol.Required(CONF_API_KEY): cv.string,
         vol.Required(CONF_STATION): cv.string,
-        vol.Required(CONF_MONITORED_CONDITIONS, default=[]): [vol.In(SENSOR_TYPES)],
+        vol.Required(CONF_MONITORED_CONDITIONS, default=[]): [vol.In(SENSOR_KEYS)],
     }
 )
 
@@ -132,44 +155,37 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
 
     weather_api = TrafikverketWeather(web_session, sensor_api)
 
-    dev = []
+    monitored_conditions = config[CONF_MONITORED_CONDITIONS]
-    for condition in config[CONF_MONITORED_CONDITIONS]:
+    entities = [
-        dev.append(
+        TrafikverketWeatherStation(
-            TrafikverketWeatherStation(
+            weather_api, sensor_name, sensor_station, description
-                weather_api, sensor_name, condition, sensor_station
-            )
         )
+        for description in SENSOR_TYPES
+        if description.key in monitored_conditions
+    ]
 
-    if dev:
+    async_add_entities(entities, True)
-        async_add_entities(dev, True)
 
 
 class TrafikverketWeatherStation(SensorEntity):
     """Representation of a Trafikverket sensor."""
 
-    def __init__(self, weather_api, name, sensor_type, sensor_station):
+    entity_description: TrafikverketSensorEntityDescription
+
+    def __init__(
+        self,
+        weather_api,
+        name,
+        sensor_station,
+        description: TrafikverketSensorEntityDescription,
+    ):
         """Initialize the sensor."""
-        self._client = name
+        self.entity_description = description
-        self._name = SENSOR_TYPES[sensor_type][0]
+        self._attr_name = f"{name} {description.name}"
-        self._type = sensor_type
-        self._state = None
-        self._unit = SENSOR_TYPES[sensor_type][1]
         self._station = sensor_station
         self._weather_api = weather_api
-        self._icon = SENSOR_TYPES[sensor_type][3]
-        self._device_class = SENSOR_TYPES[sensor_type][4]
         self._weather = None
 
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return f"{self._client} {self._name}"
-
-    @property
-    def icon(self):
-        """Icon to use in the frontend."""
-        return self._icon
-
     @property
     def extra_state_attributes(self):
         """Return the state attributes of Trafikverket Weatherstation."""
@@ -179,26 +195,13 @@ class TrafikverketWeatherStation(SensorEntity):
             ATTR_MEASURE_TIME: self._weather.measure_time,
         }
 
-    @property
-    def device_class(self):
-        """Return the device class of the sensor."""
-        return self._device_class
-
-    @property
-    def native_value(self):
-        """Return the state of the device."""
-        return self._state
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit of measurement of this entity, if any."""
-        return self._unit
-
     @Throttle(MIN_TIME_BETWEEN_UPDATES)
     async def async_update(self):
         """Get the latest data from Trafikverket and updates the states."""
         try:
             self._weather = await self._weather_api.async_get_weather(self._station)
-            self._state = getattr(self._weather, SENSOR_TYPES[self._type][2])
+            self._attr_native_value = getattr(
+                self._weather, self.entity_description.api_key
+            )
         except (asyncio.TimeoutError, aiohttp.ClientError, ValueError) as error:
             _LOGGER.error("Could not fetch weather data: %s", error)