python-kasa · rytilahti · Apr 24, 2024 · Apr 28, 2024 · Apr 28, 2024 · Apr 28, 2024
diff --git a/kasa/feature.py b/kasa/feature.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, Callable
 
@@ -14,6 +14,67 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+@dataclass
+class HassCompat:
+    """Container class for homeassistant compat.
+
+    This is used to define some metadata for homeassistant for better UX.
+    This can be directly passed to *Description constructors.
+    """
+
+    class DeviceClass(Enum):
+        """Device class for homeassistant compat.
+
+        This contains only classes used in this library.
+        """
+
+        # Sensor
+        Power = "power"  # W, kW
+        Energy = "energy"  # wH, kWh
+        Battery = "battery"  # %
+        Voltage = "voltage"  # V, mV
+        Current = "current"  # A, mA
+        Humidity = "humidity"  # %
+        Temperature = "temperature"  # ⁰C, ⁰F, K
+        Timestamp = "timestamp"  # ISO8601
+
+        # Binary sensor
+        LowBattery = "battery"
+        Connected = "connectivity"
+        # TODO: We don't want duplicate opened state, which one to use?
+        DoorOpen = "door"
+        WindowOpen = "window"
+        Overheated = "heat"
+        Wet = "moisture"
+        Problem = "problem"
+        UpdateAvailable = "update"
+
+        def __str__(self):
+            """Overridden to return only the value."""
+            return self.value
+
+    class StateClass(Enum):
+        """State class compat for homeassistant."""
+
+        Measurement = "measurement"
+        Total = "total"
+        TotalIncreasing = "total_increasing"
+
+        def __str__(self):
+            """Overridden to return only the value."""
+            return self.value
+
+    device_class: DeviceClass | None = None
+    state_class: StateClass | None = None
+    entity_registry_enabled_default: bool = True
+    entity_registry_visible_default: bool = True
+
+    def dict(self):
+        """Convert to dict ready to consume by homeassistant description classes."""
+        items = asdict(self).items()
+        return {k: v for k, v in items if v is not None}
+
+
 @dataclass
 class Feature:
     """Feature defines a generic interface for device features."""
@@ -88,6 +149,9 @@ class Category(Enum):
     #: If set, this property will be used to set *minimum_value* and *maximum_value*.
     range_getter: str | None = None
 
+    #: Homeassistant compat
+    hass_compat: HassCompat | None = None
+
     #: Identifier
     id: str | None = None
 
@@ -120,6 +184,9 @@ def __post_init__(self):
                 f" {self.type}"
             )
 
+        if self.hass_compat is None:
+            self.hass_compat = HassCompat()
+
     @property
     def value(self):
         """Return the current value."""

diff --git a/kasa/iot/iotdevice.py b/kasa/iot/iotdevice.py
@@ -25,7 +25,7 @@
 from ..deviceconfig import DeviceConfig
 from ..emeterstatus import EmeterStatus
 from ..exceptions import KasaException
-from ..feature import Feature
+from ..feature import Feature, HassCompat
 from ..protocol import BaseProtocol
 from .iotmodule import IotModule
 from .modules import Emeter
@@ -311,6 +311,7 @@ async def _initialize_features(self):
                 attribute_getter="rssi",
                 icon="mdi:signal",
                 category=Feature.Category.Debug,
+                hass_compat=HassCompat(entity_registry_enabled_default=False),
             )
         )
         if "on_time" in self._sys_info:
@@ -320,6 +321,9 @@ async def _initialize_features(self):
                     name="On since",
                     attribute_getter="on_since",
                     icon="mdi:clock",
+                    # TODO: We could enable this if we implement timezone support
+                    # hass_compat=HassCompat(device_class=
+                    # HassCompat.DeviceClass.Timestamp),
                 )
             )
 

diff --git a/kasa/iot/iotdimmer.py b/kasa/iot/iotdimmer.py
@@ -96,6 +96,7 @@ async def _initialize_features(self):
                     attribute_setter="set_brightness",
                     minimum_value=1,
                     maximum_value=100,
+                    category=Feature.Category.Primary,
                     type=Feature.Type.Number,
                 )
             )

diff --git a/kasa/iot/modules/ambientlight.py b/kasa/iot/modules/ambientlight.py
@@ -26,6 +26,7 @@ def __init__(self, device, module):
                 icon="mdi:brightness-percent",
                 attribute_getter="ambientlight_brightness",
                 type=Feature.Type.Sensor,
+                unit="%",
             )
         )
 

diff --git a/kasa/iot/modules/cloud.py b/kasa/iot/modules/cloud.py
@@ -5,7 +5,7 @@
 except ImportError:
     from pydantic import BaseModel
 
-from ...feature import Feature
+from ...feature import Feature, HassCompat
 from ..iotmodule import IotModule
 
 
@@ -37,6 +37,8 @@ def __init__(self, device, module):
                 icon="mdi:cloud",
                 attribute_getter="is_connected",
                 type=Feature.Type.BinarySensor,
+                category=Feature.Category.Debug,
+                hass_compat=HassCompat(device_class=HassCompat.DeviceClass.Connected),
             )
         )
 

diff --git a/kasa/iot/modules/emeter.py b/kasa/iot/modules/emeter.py
@@ -6,7 +6,7 @@
 
 from ... import Device
 from ...emeterstatus import EmeterStatus
-from ...feature import Feature
+from ...feature import Feature, HassCompat
 from .usage import Usage
 
 
@@ -23,6 +23,10 @@ def __init__(self, device: Device, module: str):
                 container=self,
                 unit="W",
                 id="current_power_w",  # for homeassistant backwards compat
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Power,
+                    state_class=HassCompat.StateClass.Measurement,
+                ),
             )
         )
         self._add_feature(
@@ -33,6 +37,11 @@ def __init__(self, device: Device, module: str):
                 container=self,
                 unit="kWh",
                 id="today_energy_kwh",  # for homeassistant backwards compat
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Energy,
+                    state_class=HassCompat.StateClass.TotalIncreasing,
+                    entity_registry_visible_default=False,
+                ),
             )
         )
         self._add_feature(
@@ -42,6 +51,11 @@ def __init__(self, device: Device, module: str):
                 attribute_getter="emeter_this_month",
                 container=self,
                 unit="kWh",
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Energy,
+                    state_class=HassCompat.StateClass.TotalIncreasing,
+                    entity_registry_visible_default=False,
+                ),
             )
         )
         self._add_feature(
@@ -52,6 +66,11 @@ def __init__(self, device: Device, module: str):
                 container=self,
                 unit="kWh",
                 id="total_energy_kwh",  # for homeassistant backwards compat
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Energy,
+                    state_class=HassCompat.StateClass.TotalIncreasing,
+                    entity_registry_visible_default=False,
+                ),
             )
         )
         self._add_feature(
@@ -62,6 +81,10 @@ def __init__(self, device: Device, module: str):
                 container=self,
                 unit="V",
                 id="voltage",  # for homeassistant backwards compat
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Voltage,
+                    state_class=HassCompat.StateClass.Measurement,
+                ),
             )
         )
         self._add_feature(
@@ -72,6 +95,10 @@ def __init__(self, device: Device, module: str):
                 container=self,
                 unit="A",
                 id="current_a",  # for homeassistant backwards compat
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Current,
+                    state_class=HassCompat.StateClass.Measurement,
+                ),
             )
         )
 

diff --git a/kasa/smart/modules/__init__.py b/kasa/smart/modules/__init__.py
@@ -12,6 +12,7 @@
 from .energymodule import EnergyModule
 from .fanmodule import FanModule
 from .firmware import Firmware
+from .frostprotection import FrostProtectionModule
 from .humidity import HumiditySensor
 from .ledmodule import LedModule
 from .lighttransitionmodule import LightTransitionModule
@@ -40,4 +41,5 @@
     "LightTransitionModule",
     "ColorTemperatureModule",
     "ColorModule",
+    "FrostProtectionModule",
 ]
diff --git a/kasa/smart/modules/autooffmodule.py b/kasa/smart/modules/autooffmodule.py
@@ -42,7 +42,11 @@ def __init__(self, device: SmartDevice, module: str):
         )
         self._add_feature(
             Feature(
-                device, "Auto off at", container=self, attribute_getter="auto_off_at"
+                device,
+                "Auto off at",
+                container=self,
+                attribute_getter="auto_off_at",
+                category=Feature.Category.Debug,
             )
         )
 

diff --git a/kasa/smart/modules/battery.py b/kasa/smart/modules/battery.py
@@ -4,7 +4,7 @@
 
 from typing import TYPE_CHECKING
 
-from ...feature import Feature
+from ...feature import Feature, HassCompat
 from ..smartmodule import SmartModule
 
 if TYPE_CHECKING:
@@ -26,6 +26,9 @@ def __init__(self, device: SmartDevice, module: str):
                 container=self,
                 attribute_getter="battery",
                 icon="mdi:battery",
+                unit="%",
+                category=Feature.Category.Debug,
+                hass_compat=HassCompat(device_class=HassCompat.DeviceClass.Battery),
             )
         )
         self._add_feature(
@@ -36,6 +39,8 @@ def __init__(self, device: SmartDevice, module: str):
                 attribute_getter="battery_low",
                 icon="mdi:alert",
                 type=Feature.Type.BinarySensor,
+                category=Feature.Category.Debug,
+                hass_compat=HassCompat(device_class=HassCompat.DeviceClass.LowBattery),
             )
         )
 

diff --git a/kasa/smart/modules/cloudmodule.py b/kasa/smart/modules/cloudmodule.py
@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING
 
 from ...exceptions import SmartErrorCode
-from ...feature import Feature
+from ...feature import Feature, HassCompat
 from ..smartmodule import SmartModule
 
 if TYPE_CHECKING:
@@ -29,6 +29,8 @@ def __init__(self, device: SmartDevice, module: str):
                 attribute_getter="is_connected",
                 icon="mdi:cloud",
                 type=Feature.Type.BinarySensor,
+                category=Feature.Category.Debug,
+                hass_compat=HassCompat(device_class=HassCompat.DeviceClass.Connected),
             )
         )
 

diff --git a/kasa/smart/modules/energymodule.py b/kasa/smart/modules/energymodule.py
@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING
 
 from ...emeterstatus import EmeterStatus
-from ...feature import Feature
+from ...feature import Feature, HassCompat
 from ..smartmodule import SmartModule
 
 if TYPE_CHECKING:
@@ -26,6 +26,10 @@ def __init__(self, device: SmartDevice, module: str):
                 attribute_getter="current_power",
                 container=self,
                 unit="W",
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Power,
+                    state_class=HassCompat.StateClass.Measurement,
+                ),
             )
         )
         self._add_feature(
@@ -35,6 +39,10 @@ def __init__(self, device: SmartDevice, module: str):
                 attribute_getter="emeter_today",
                 container=self,
                 unit="Wh",
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Energy,
+                    state_class=HassCompat.StateClass.TotalIncreasing,
+                ),
             )
         )
         self._add_feature(
@@ -44,6 +52,10 @@ def __init__(self, device: SmartDevice, module: str):
                 attribute_getter="emeter_this_month",
                 container=self,
                 unit="Wh",
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.Energy,
+                    state_class=HassCompat.StateClass.TotalIncreasing,
+                ),
             )
         )
 

diff --git a/kasa/smart/modules/firmware.py b/kasa/smart/modules/firmware.py
@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 from ...exceptions import SmartErrorCode
-from ...feature import Feature
+from ...feature import Feature, HassCompat
 from ..smartmodule import SmartModule
 
 try:
@@ -69,6 +69,10 @@ def __init__(self, device: SmartDevice, module: str):
                 container=self,
                 attribute_getter="update_available",
                 type=Feature.Type.BinarySensor,
+                category=Feature.Category.Debug,
+                hass_compat=HassCompat(
+                    device_class=HassCompat.DeviceClass.UpdateAvailable
+                ),
             )
         )