Initial Commit

custom_components/asusrouter/entity.py

@@ -0,0 +1,192 @@
+"""AsusRouter entity module."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from asusrouter.modules.homeassistant import convert_to_ha_state_bool
+from asusrouter.modules.state import AsusState
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
+
+from .const import ASUSROUTER, COORDINATOR, DOMAIN
+from .dataclass import (
+    ARBinaryDescription,
+    AREntityDescription,
+    ARSwitchDescription,
+)
+from .helpers import to_unique_id
+from .router import ARDevice
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_ar_entry(  # noqa: PLR0913
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+    sensors: list[AREntityDescription],
+    sensor_class: type[AREntity],
+    hide: list[str] | None = None,
+) -> None:
+    """Set up AsusRouter entities."""
+
+    router: ARDevice = hass.data[DOMAIN][config_entry.entry_id][ASUSROUTER]
+    entities = []
+
+    if not hide:
+        hide = []
+
+    for sensor_data in router.sensor_coordinator.values():
+        coordinator = sensor_data[COORDINATOR]
+        for sensor_description in sensors:
+            try:
+                sensor_type = sensor_description.key_group
+
+                # Make sure extra state attributes are dict
+                if not sensor_description.extra_state_attributes:
+                    sensor_description.extra_state_attributes = {}
+
+                if (
+                    sensor_type in sensor_data
+                    and sensor_description.key in sensor_data[sensor_type]
+                ):
+                    # Hide protected values
+                    sensor_description.extra_state_attributes = {
+                        key: value
+                        for key, value in sensor_description.extra_state_attributes.items()  # noqa: E501
+                        if value not in hide
+                    }
+
+                    entities.append(
+                        sensor_class(coordinator, router, sensor_description)
+                    )
+            except Exception as ex:  # noqa: BLE001
+                _LOGGER.warning(
+                    "Got an exception when creating entities: %s. "
+                    "Please, report this. Sensor description: %s",
+                    ex,
+                    sensor_description,
+                )
+
+    async_add_entities(entities)
+
+
+class AREntity(CoordinatorEntity):
+    """AsusRouter entity."""
+
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        router: ARDevice,
+        description: AREntityDescription,
+    ) -> None:
+        """Initialize AsusRouter entity."""
+
+        super().__init__(coordinator)
+        self.router = router
+        self.api = router.bridge.api
+        self.coordinator = coordinator
+
+        self._attr_name = f"{router._conf_name} {description.name}"
+        self._attr_unique_id = to_unique_id(f"{router.mac}_{description.name}")
+        self._attr_device_info = router.device_info
+        self._attr_capability_attributes = description.capabilities
+
+    @property
+    def extra_state_attributes(self) -> dict[str, Any]:
+        """Return extra state attributes."""
+
+        # Check if description is of the needed class
+        if not isinstance(self.entity_description, AREntityDescription):
+            return {}
+
+        description = self.entity_description
+        _attributes = description.extra_state_attributes
+        if not _attributes:
+            return {}
+
+        # Mask attributes
+        attrs_to_mask = self.router.sensor_filters.get(
+            (description.key_group, description.key),
+            [],
+        )
+        for attr in attrs_to_mask:
+            if attr in _attributes:
+                _attributes.pop(attr, None)
+
+        attributes = {}
+
+        for attr in _attributes:
+            if attr in self.coordinator.data:
+                attributes[_attributes[attr]] = self.coordinator.data[attr]
+
+        return dict(sorted(attributes.items())) or {}
+
+
+class ARBinaryEntity(AREntity):
+    """AsusRouter binary entity."""
+
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        router: ARDevice,
+        description: AREntityDescription,
+    ) -> None:
+        """Initialize AsusRouter binary entity."""
+
+        super().__init__(coordinator, router, description)
+        if isinstance(description, ARBinaryDescription | ARSwitchDescription):
+            self._icon_onoff = bool(
+                description.icon_on and description.icon_off
+            )
+
+    @property
+    def is_on(self) -> bool | None:
+        """Get the state."""
+
+        return convert_to_ha_state_bool(
+            self.coordinator.data.get(self.entity_description.key)
+        )
+
+    @property
+    def icon(self) -> str | None:
+        """Get the icon."""
+
+        if (
+            isinstance(
+                self.entity_description,
+                ARBinaryDescription | ARSwitchDescription,
+            )
+            and self._icon_onoff
+        ):
+            if self.is_on:
+                return self.entity_description.icon_on
+            return self.entity_description.icon_off
+        if self.entity_description.icon:
+            return self.entity_description.icon
+        return None
+
+    async def _set_state(
+        self,
+        state: AsusState,
+        expect_modify: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Set switch state."""
+        try:
+            _LOGGER.debug("Setting state to %s", state)
+            result = await self.api.async_set_state(
+                state=state, expect_modify=expect_modify, **kwargs
+            )
+            await self.coordinator.async_request_refresh()
+            if not result:
+                _LOGGER.debug("State was not set!")
+        except Exception as ex:  # noqa: BLE001
+            _LOGGER.error("Unable to set state with an exception: %s", ex)