Use EntityDescription - notion (#55120)

homeassistant/components/notion/sensor.py

@@ -1,17 +1,21 @@
 """Support for Notion sensors."""
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import DEVICE_CLASS_TEMPERATURE, TEMP_CELSIUS
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from . import NotionEntity
 from .const import DATA_COORDINATOR, DOMAIN, LOGGER, SENSOR_TEMPERATURE
 
-SENSOR_TYPES = {
-    SENSOR_TEMPERATURE: ("Temperature", DEVICE_CLASS_TEMPERATURE, TEMP_CELSIUS)
-}
+SENSOR_DESCRIPTIONS = (
+    SensorEntityDescription(
+        key=SENSOR_TEMPERATURE,
+        name="Temperature",
+        device_class=DEVICE_CLASS_TEMPERATURE,
+        native_unit_of_measurement=TEMP_CELSIUS,
+    ),
+)
 
 
 async def async_setup_entry(
@@ -20,51 +24,27 @@ async def async_setup_entry(
     """Set up Notion sensors based on a config entry."""
     coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id]
 
-    sensor_list = []
-    for task_id, task in coordinator.data["tasks"].items():
-        if task["task_type"] not in SENSOR_TYPES:
-            continue
-
-        name, device_class, unit = SENSOR_TYPES[task["task_type"]]
-        sensor = coordinator.data["sensors"][task["sensor_id"]]
-
-        sensor_list.append(
+    async_add_entities(
+        [
             NotionSensor(
                 coordinator,
                 task_id,
                 sensor["id"],
                 sensor["bridge"]["id"],
                 sensor["system_id"],
-                name,
-                device_class,
-                unit,
+                description,
             )
-        )
-
-    async_add_entities(sensor_list)
+            for task_id, task in coordinator.data["tasks"].items()
+            for description in SENSOR_DESCRIPTIONS
+            if description.key == task["task_type"]
+            and (sensor := coordinator.data["sensors"][task["sensor_id"]])
+        ]
+    )
 
 
 class NotionSensor(NotionEntity, SensorEntity):
     """Define a Notion sensor."""
 
-    def __init__(
-        self,
-        coordinator: DataUpdateCoordinator,
-        task_id: str,
-        sensor_id: str,
-        bridge_id: str,
-        system_id: str,
-        name: str,
-        device_class: str,
-        unit: str,
-    ) -> None:
-        """Initialize the entity."""
-        super().__init__(
-            coordinator, task_id, sensor_id, bridge_id, system_id, name, device_class
-        )
-
-        self._attr_native_unit_of_measurement = unit
-
     @callback
     def _async_update_from_latest_data(self) -> None:
         """Fetch new state data for the sensor."""