Add Base class for entities in PyLoad integration (#120563)

* Add Base class for entities

* Remove constructors

homeassistant/components/pyload/button.py

@@ -11,13 +11,10 @@ from pyloadapi.api import PyLoadAPI
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import PyLoadConfigEntry
-from .const import DOMAIN, MANUFACTURER, SERVICE_NAME
-from .coordinator import PyLoadCoordinator
+from .entity import BasePyLoadEntity
 
 
 @dataclass(kw_only=True, frozen=True)
@@ -76,32 +73,11 @@ async def async_setup_entry(
     )
 
 
-class PyLoadBinarySensor(CoordinatorEntity[PyLoadCoordinator], ButtonEntity):
+class PyLoadBinarySensor(BasePyLoadEntity, ButtonEntity):
     """Representation of a pyLoad button."""
 
-    _attr_has_entity_name = True
     entity_description: PyLoadButtonEntityDescription
 
-    def __init__(
-        self,
-        coordinator: PyLoadCoordinator,
-        entity_description: PyLoadButtonEntityDescription,
-    ) -> None:
-        """Initialize the button."""
-        super().__init__(coordinator)
-        self._attr_unique_id = (
-            f"{coordinator.config_entry.entry_id}_{entity_description.key}"
-        )
-        self.entity_description = entity_description
-        self._attr_device_info = DeviceInfo(
-            entry_type=DeviceEntryType.SERVICE,
-            manufacturer=MANUFACTURER,
-            model=SERVICE_NAME,
-            configuration_url=coordinator.pyload.api_url,
-            identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
-            sw_version=coordinator.version,
-        )
-
     async def async_press(self) -> None:
         """Handle the button press."""
         await self.entity_description.press_fn(self.coordinator.pyload)

homeassistant/components/pyload/entity.py

@@ -0,0 +1,37 @@
+"""Base entity for pyLoad."""
+
+from __future__ import annotations
+
+from homeassistant.components.button import EntityDescription
+from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import DOMAIN, MANUFACTURER, SERVICE_NAME
+from .coordinator import PyLoadCoordinator
+
+
+class BasePyLoadEntity(CoordinatorEntity[PyLoadCoordinator]):
+    """BaseEntity for pyLoad."""
+
+    _attr_has_entity_name = True
+
+    def __init__(
+        self,
+        coordinator: PyLoadCoordinator,
+        entity_description: EntityDescription,
+    ) -> None:
+        """Initialize the Entity."""
+        super().__init__(coordinator)
+
+        self._attr_unique_id = (
+            f"{coordinator.config_entry.entry_id}_{entity_description.key}"
+        )
+        self.entity_description = entity_description
+        self._attr_device_info = DeviceInfo(
+            entry_type=DeviceEntryType.SERVICE,
+            manufacturer=MANUFACTURER,
+            model=SERVICE_NAME,
+            configuration_url=coordinator.pyload.api_url,
+            identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
+            sw_version=coordinator.version,
+        )

homeassistant/components/pyload/sensor.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
+from dataclasses import dataclass
 from enum import StrEnum
 
 import voluptuous as vol
@@ -28,11 +30,9 @@ from homeassistant.const import (
 from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 import homeassistant.helpers.config_validation as cv
-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, StateType
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import PyLoadConfigEntry
 from .const import (
@@ -41,11 +41,10 @@ from .const import (
     DEFAULT_PORT,
     DOMAIN,
     ISSUE_PLACEHOLDER,
-    MANUFACTURER,
-    SERVICE_NAME,
     UNIT_DOWNLOADS,
 )
-from .coordinator import PyLoadCoordinator
+from .coordinator import pyLoadData
+from .entity import BasePyLoadEntity
 
 
 class PyLoadSensorEntity(StrEnum):
@@ -58,40 +57,52 @@ class PyLoadSensorEntity(StrEnum):
     TOTAL = "total"
 
 
-SENSOR_DESCRIPTIONS: tuple[SensorEntityDescription, ...] = (
-    SensorEntityDescription(
+@dataclass(kw_only=True, frozen=True)
+class PyLoadSensorEntityDescription(SensorEntityDescription):
+    """Describes pyLoad switch entity."""
+
+    value_fn: Callable[[pyLoadData], StateType]
+
+
+SENSOR_DESCRIPTIONS: tuple[PyLoadSensorEntityDescription, ...] = (
+    PyLoadSensorEntityDescription(
         key=PyLoadSensorEntity.SPEED,
         translation_key=PyLoadSensorEntity.SPEED,
         device_class=SensorDeviceClass.DATA_RATE,
         native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND,
         suggested_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND,
         suggested_display_precision=1,
+        value_fn=lambda data: data.speed,
     ),
-    SensorEntityDescription(
+    PyLoadSensorEntityDescription(
         key=PyLoadSensorEntity.ACTIVE,
         translation_key=PyLoadSensorEntity.ACTIVE,
         native_unit_of_measurement=UNIT_DOWNLOADS,
         state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda data: data.active,
     ),
-    SensorEntityDescription(
+    PyLoadSensorEntityDescription(
         key=PyLoadSensorEntity.QUEUE,
         translation_key=PyLoadSensorEntity.QUEUE,
         native_unit_of_measurement=UNIT_DOWNLOADS,
         state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda data: data.queue,
     ),
-    SensorEntityDescription(
+    PyLoadSensorEntityDescription(
         key=PyLoadSensorEntity.TOTAL,
         translation_key=PyLoadSensorEntity.TOTAL,
         native_unit_of_measurement=UNIT_DOWNLOADS,
         state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda data: data.total,
     ),
-    SensorEntityDescription(
+    PyLoadSensorEntityDescription(
         key=PyLoadSensorEntity.FREE_SPACE,
         translation_key=PyLoadSensorEntity.FREE_SPACE,
         device_class=SensorDeviceClass.DATA_SIZE,
         native_unit_of_measurement=UnitOfInformation.BYTES,
         suggested_unit_of_measurement=UnitOfInformation.GIBIBYTES,
         suggested_display_precision=1,
+        value_fn=lambda data: data.free_space,
     ),
 )
 
@@ -173,32 +184,12 @@ async def async_setup_entry(
     )
 
 
-class PyLoadSensor(CoordinatorEntity[PyLoadCoordinator], SensorEntity):
+class PyLoadSensor(BasePyLoadEntity, SensorEntity):
     """Representation of a pyLoad sensor."""
 
-    _attr_has_entity_name = True
-
-    def __init__(
-        self,
-        coordinator: PyLoadCoordinator,
-        entity_description: SensorEntityDescription,
-    ) -> None:
-        """Initialize a new pyLoad sensor."""
-        super().__init__(coordinator)
-        self._attr_unique_id = (
-            f"{coordinator.config_entry.entry_id}_{entity_description.key}"
-        )
-        self.entity_description = entity_description
-        self._attr_device_info = DeviceInfo(
-            entry_type=DeviceEntryType.SERVICE,
-            manufacturer=MANUFACTURER,
-            model=SERVICE_NAME,
-            configuration_url=coordinator.pyload.api_url,
-            identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
-            sw_version=coordinator.version,
-        )
+    entity_description: PyLoadSensorEntityDescription
 
     @property
     def native_value(self) -> StateType:
         """Return the state of the sensor."""
-        return getattr(self.coordinator.data, self.entity_description.key)
+        return self.entity_description.value_fn(self.coordinator.data)

homeassistant/components/pyload/switch.py

@@ -15,13 +15,10 @@ from homeassistant.components.switch import (
     SwitchEntityDescription,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import PyLoadConfigEntry
-from .const import DOMAIN, MANUFACTURER, SERVICE_NAME
-from .coordinator import PyLoadCoordinator
+from .entity import BasePyLoadEntity
 
 
 class PyLoadSwitch(StrEnum):
@@ -75,32 +72,11 @@ async def async_setup_entry(
     )
 
 
-class PyLoadSwitchEntity(CoordinatorEntity[PyLoadCoordinator], SwitchEntity):
+class PyLoadSwitchEntity(BasePyLoadEntity, SwitchEntity):
     """Representation of a pyLoad sensor."""
 
-    _attr_has_entity_name = True
     entity_description: PyLoadSwitchEntityDescription
 
-    def __init__(
-        self,
-        coordinator: PyLoadCoordinator,
-        entity_description: PyLoadSwitchEntityDescription,
-    ) -> None:
-        """Initialize the sensor."""
-        super().__init__(coordinator)
-        self._attr_unique_id = (
-            f"{coordinator.config_entry.entry_id}_{entity_description.key}"
-        )
-        self.entity_description = entity_description
-        self._attr_device_info = DeviceInfo(
-            entry_type=DeviceEntryType.SERVICE,
-            manufacturer=MANUFACTURER,
-            model=SERVICE_NAME,
-            configuration_url=coordinator.pyload.api_url,
-            identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
-            sw_version=coordinator.version,
-        )
-
     @property
     def is_on(self) -> bool | None:
         """Return the state of the device."""