Add generic interface for accessing device features (#741)

This adds a generic interface for all device classes to introspect available device features,
that is necessary to make it easier to support a wide variety of supported devices with different set of features.
This will allow constructing generic interfaces (e.g., in homeassistant) that fetch and change these features without hard-coding the API calls.

`Device.features()` now returns a mapping of `<identifier, Feature>` where the `Feature` contains all necessary information (like the name, the icon, a way to get and change the setting) to present and change the defined feature through its interface.

kasa/__init__.py

@@ -33,6 +33,7 @@ from kasa.exceptions import (
     TimeoutException,
     UnsupportedDeviceException,
 )
+from kasa.feature import Feature, FeatureType
 from kasa.iot.iotbulb import BulbPreset, TurnOnBehavior, TurnOnBehaviors
 from kasa.iotprotocol import (
     IotProtocol,
@@ -54,6 +55,8 @@ __all__ = [
     "TurnOnBehaviors",
     "TurnOnBehavior",
     "DeviceType",
+    "Feature",
+    "FeatureType",
     "EmeterStatus",
     "Device",
     "Bulb",

kasa/cli.py

@@ -102,6 +102,7 @@ class ExceptionHandlerGroup(click.Group):
             asyncio.get_event_loop().run_until_complete(self.main(*args, **kwargs))
         except Exception as ex:
             echo(f"Got error: {ex!r}")
+            raise
 
 
 def json_formatter_cb(result, **kwargs):
@@ -578,6 +579,10 @@ async def state(ctx, dev: Device):
         else:
             echo(f"\t{info_name}: {info_data}")
 
+    echo("\n\t[bold]== Features == [/bold]")
+    for id_, feature in dev.features.items():
+        echo(f"\t{feature.name} ({id_}): {feature.value}")
+
     if dev.has_emeter:
         echo("\n\t[bold]== Current State ==[/bold]")
         emeter_status = dev.emeter_realtime
@@ -594,8 +599,6 @@ async def state(ctx, dev: Device):
         echo("\n\t[bold]== Verbose information ==[/bold]")
         echo(f"\tCredentials hash:  {dev.credentials_hash}")
         echo(f"\tDevice ID:         {dev.device_id}")
-        for feature in dev.features:
-            echo(f"\tFeature:           {feature}")
         echo()
         _echo_discovery_info(dev._discovery_info)
     return dev.internal_state
@@ -1115,5 +1118,37 @@ async def shell(dev: Device):
         loop.stop()
 
 
+@cli.command(name="feature")
+@click.argument("name", required=False)
+@click.argument("value", required=False)
+@pass_dev
+async def feature(dev, name: str, value):
+    """Access and modify features.
+
+    If no *name* is given, lists available features and their values.
+    If only *name* is given, the value of named feature is returned.
+    If both *name* and *value* are set, the described setting is changed.
+    """
+    if not name:
+        echo("[bold]== Features ==[/bold]")
+        for name, feat in dev.features.items():
+            echo(f"{feat.name} ({name}): {feat.value}")
+        return
+
+    if name not in dev.features:
+        echo(f"No feature by name {name}")
+        return
+
+    feat = dev.features[name]
+
+    if value is None:
+        echo(f"{feat.name} ({name}): {feat.value}")
+        return feat.value
+
+    echo(f"Setting {name} to {value}")
+    value = ast.literal_eval(value)
+    return await dev.features[name].set_value(value)
+
+
 if __name__ == "__main__":
     cli()

kasa/device.py

@@ -3,13 +3,14 @@ import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Sequence, Set, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 from .credentials import Credentials
 from .device_type import DeviceType
 from .deviceconfig import DeviceConfig
 from .emeterstatus import EmeterStatus
 from .exceptions import SmartDeviceException
+from .feature import Feature
 from .iotprotocol import IotProtocol
 from .protocol import BaseProtocol
 from .xortransport import XorTransport
@@ -69,6 +70,7 @@ class Device(ABC):
         self._discovery_info: Optional[Dict[str, Any]] = None
 
         self.modules: Dict[str, Any] = {}
+        self._features: Dict[str, Feature] = {}
 
     @staticmethod
     async def connect(
@@ -296,9 +298,16 @@ class Device(ABC):
         """Return the key state information."""
 
     @property
-    @abstractmethod
+    def features(self) -> Dict[str, Feature]:
-    def features(self) -> Set[str]:
         """Return the list of supported features."""
+        return self._features
+
+    def _add_feature(self, feature: Feature):
+        """Add a new feature to the device."""
+        desc_name = feature.name.lower().replace(" ", "_")
+        if desc_name in self._features:
+            raise SmartDeviceException("Duplicate feature name %s" % desc_name)
+        self._features[desc_name] = feature
 
     @property
     @abstractmethod

kasa/feature.py

@@ -0,0 +1,50 @@
+"""Generic interface for defining device features."""
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+
+if TYPE_CHECKING:
+    from .device import Device
+
+
+class FeatureType(Enum):
+    """Type to help decide how to present the feature."""
+
+    Sensor = auto()
+    BinarySensor = auto()
+    Switch = auto()
+    Button = auto()
+
+
+@dataclass
+class Feature:
+    """Feature defines a generic interface for device features."""
+
+    #: Device instance required for getting and setting values
+    device: "Device"
+    #: User-friendly short description
+    name: str
+    #: Name of the property that allows accessing the value
+    attribute_getter: Union[str, Callable]
+    #: Name of the method that allows changing the value
+    attribute_setter: Optional[str] = None
+    #: Container storing the data, this overrides 'device' for getters
+    container: Any = None
+    #: Icon suggestion
+    icon: Optional[str] = None
+    #: Type of the feature
+    type: FeatureType = FeatureType.Sensor
+
+    @property
+    def value(self):
+        """Return the current value."""
+        container = self.container if self.container is not None else self.device
+        if isinstance(self.attribute_getter, Callable):
+            return self.attribute_getter(container)
+        return getattr(container, self.attribute_getter)
+
+    async def set_value(self, value):
+        """Set the value."""
+        if self.attribute_setter is None:
+            raise ValueError("Tried to set read-only feature.")
+        return await getattr(self.device, self.attribute_setter)(value)

kasa/iot/iotdevice.py

@@ -22,6 +22,7 @@ from ..device import Device, WifiNetwork
 from ..deviceconfig import DeviceConfig
 from ..emeterstatus import EmeterStatus
 from ..exceptions import SmartDeviceException
+from ..feature import Feature
 from ..protocol import BaseProtocol
 from .modules import Emeter, IotModule
 
@@ -184,8 +185,9 @@ class IotDevice(Device):
         super().__init__(host=host, config=config, protocol=protocol)
 
         self._sys_info: Any = None  # TODO: this is here to avoid changing tests
-        self._features: Set[str] = set()
         self._children: Sequence["IotDevice"] = []
+        self._supported_modules: Optional[Dict[str, IotModule]] = None
+        self._legacy_features: Set[str] = set()
 
     @property
     def children(self) -> Sequence["IotDevice"]:
@@ -260,7 +262,7 @@ class IotDevice(Device):
 
     @property  # type: ignore
     @requires_update
-    def features(self) -> Set[str]:
+    def features(self) -> Dict[str, Feature]:
         """Return a set of features that the device supports."""
         return self._features
 
@@ -276,7 +278,7 @@ class IotDevice(Device):
     @requires_update
     def has_emeter(self) -> bool:
         """Return True if device has an energy meter."""
-        return "ENE" in self.features
+        return "ENE" in self._legacy_features
 
     async def get_sys_info(self) -> Dict[str, Any]:
         """Retrieve system information."""
@@ -299,9 +301,28 @@ class IotDevice(Device):
             self._last_update = response
             self._set_sys_info(response["system"]["get_sysinfo"])
 
+        if not self._features:
+            await self._initialize_features()
+
         await self._modular_update(req)
         self._set_sys_info(self._last_update["system"]["get_sysinfo"])
 
+    async def _initialize_features(self):
+        self._add_feature(
+            Feature(
+                device=self, name="RSSI", attribute_getter="rssi", icon="mdi:signal"
+            )
+        )
+        if "on_time" in self._sys_info:
+            self._add_feature(
+                Feature(
+                    device=self,
+                    name="On since",
+                    attribute_getter="on_since",
+                    icon="mdi:clock",
+                )
+            )
+
     async def _modular_update(self, req: dict) -> None:
         """Execute an update query."""
         if self.has_emeter:
@@ -310,6 +331,18 @@ class IotDevice(Device):
             )
             self.add_module("emeter", Emeter(self, self.emeter_type))
 
+        # TODO: perhaps modules should not have unsupported modules,
+        #  making separate handling for this unnecessary
+        if self._supported_modules is None:
+            supported = {}
+            for module in self.modules.values():
+                if module.is_supported:
+                    supported[module._module] = module
+                for module_feat in module._module_features.values():
+                    self._add_feature(module_feat)
+
+            self._supported_modules = supported
+
         request_list = []
         est_response_size = 1024 if "system" in req else 0
         for module in self.modules.values():
@@ -357,9 +390,7 @@ class IotDevice(Device):
         """Set sys_info."""
         self._sys_info = sys_info
         if features := sys_info.get("feature"):
-            self._features = _parse_features(features)
+            self._legacy_features = _parse_features(features)
-        else:
-            self._features = set()
 
     @property  # type: ignore
     @requires_update

kasa/iot/iotplug.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, Optional
 
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
+from ..feature import Feature, FeatureType
 from ..protocol import BaseProtocol
 from .iotdevice import IotDevice, requires_update
 from .modules import Antitheft, Cloud, Schedule, Time, Usage
@@ -56,6 +57,17 @@ class IotPlug(IotDevice):
         self.add_module("time", Time(self, "time"))
         self.add_module("cloud", Cloud(self, "cnCloud"))
 
+        self._add_feature(
+            Feature(
+                device=self,
+                name="LED",
+                icon="mdi:led-{state}",
+                attribute_getter="led",
+                attribute_setter="set_led",
+                type=FeatureType.Switch,
+            )
+        )
+
     @property  # type: ignore
     @requires_update
     def is_on(self) -> bool:
@@ -88,5 +100,4 @@ class IotPlug(IotDevice):
     @requires_update
     def state_information(self) -> Dict[str, Any]:
         """Return switch-specific state information."""
-        info = {"LED state": self.led, "On since": self.on_since}
+        return {}
-        return info

kasa/iot/modules/cloud.py

@@ -4,6 +4,7 @@ try:
 except ImportError:
     from pydantic import BaseModel
 
+from ...feature import Feature, FeatureType
 from .module import IotModule
 
 
@@ -25,6 +26,24 @@ class CloudInfo(BaseModel):
 class Cloud(IotModule):
     """Module implementing support for cloud services."""
 
+    def __init__(self, device, module):
+        super().__init__(device, module)
+        self._add_feature(
+            Feature(
+                device=device,
+                container=self,
+                name="Cloud connection",
+                icon="mdi:cloud",
+                attribute_getter="is_connected",
+                type=FeatureType.BinarySensor,
+            )
+        )
+
+    @property
+    def is_connected(self) -> bool:
+        """Return true if device is connected to the cloud."""
+        return self.info.binded
+
     def query(self):
         """Request cloud connectivity info."""
         return self.query_for_command("get_info")

kasa/iot/modules/module.py

@@ -2,9 +2,10 @@
 import collections
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Dict
 
 from ...exceptions import SmartDeviceException
+from ...feature import Feature
 
 if TYPE_CHECKING:
     from kasa.iot import IotDevice
@@ -34,6 +35,14 @@ class IotModule(ABC):
     def __init__(self, device: "IotDevice", module: str):
         self._device = device
         self._module = module
+        self._module_features: Dict[str, Feature] = {}
+
+    def _add_feature(self, feature: Feature):
+        """Add module feature."""
+        feature_name = f"{self._module}_{feature.name}"
+        if feature_name in self._module_features:
+            raise SmartDeviceException("Duplicate name detected %s" % feature_name)
+        self._module_features[feature_name] = feature
 
     @abstractmethod
     def query(self):

kasa/smart/smartdevice.py

@@ -2,7 +2,7 @@
 import base64
 import logging
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Set, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, cast
 
 from ..aestransport import AesTransport
 from ..device import Device, WifiNetwork
@@ -10,6 +10,7 @@ from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
 from ..emeterstatus import EmeterStatus
 from ..exceptions import AuthenticationException, SmartDeviceException
+from ..feature import Feature, FeatureType
 from ..smartprotocol import SmartProtocol
 
 _LOGGER = logging.getLogger(__name__)
@@ -124,6 +125,11 @@ class SmartDevice(Device):
             for info in child_info["child_device_list"]:
                 self._children[info["device_id"]].update_internal_state(info)
 
+        # We can first initialize the features after the first update.
+        # We make here an assumption that every device has at least a single feature.
+        if not self._features:
+            await self._initialize_features()
+
         _LOGGER.debug("Got an update: %s", self._last_update)
 
     async def _initialize_modules(self):
@@ -131,6 +137,51 @@ class SmartDevice(Device):
         if "energy_monitoring" in self._components:
             self.emeter_type = "emeter"
 
+    async def _initialize_features(self):
+        """Initialize device features."""
+        self._add_feature(
+            Feature(
+                self,
+                "Signal Level",
+                attribute_getter=lambda x: x._info["signal_level"],
+                icon="mdi:signal",
+            )
+        )
+        self._add_feature(
+            Feature(
+                self,
+                "RSSI",
+                attribute_getter=lambda x: x._info["rssi"],
+                icon="mdi:signal",
+            )
+        )
+        self._add_feature(
+            Feature(device=self, name="SSID", attribute_getter="ssid", icon="mdi:wifi")
+        )
+
+        if "overheated" in self._info:
+            self._add_feature(
+                Feature(
+                    self,
+                    "Overheated",
+                    attribute_getter=lambda x: x._info["overheated"],
+                    icon="mdi:heat-wave",
+                    type=FeatureType.BinarySensor,
+                )
+            )
+
+        # We check for the key available, and not for the property truthiness,
+        # as the value is falsy when the device is off.
+        if "on_time" in self._info:
+            self._add_feature(
+                Feature(
+                    device=self,
+                    name="On since",
+                    attribute_getter="on_since",
+                    icon="mdi:clock",
+                )
+            )
+
     @property
     def sys_info(self) -> Dict[str, Any]:
         """Returns the device info."""
@@ -221,23 +272,21 @@ class SmartDevice(Device):
         return res
 
     @property
-    def state_information(self) -> Dict[str, Any]:
+    def ssid(self) -> str:
-        """Return the key state information."""
+        """Return ssid of the connected wifi ap."""
         ssid = self._info.get("ssid")
         ssid = base64.b64decode(ssid).decode() if ssid else "No SSID"
+        return ssid
 
+    @property
+    def state_information(self) -> Dict[str, Any]:
+        """Return the key state information."""
         return {
             "overheated": self._info.get("overheated"),
             "signal_level": self._info.get("signal_level"),
-            "SSID": ssid,
+            "SSID": self.ssid,
         }
 
-    @property
-    def features(self) -> Set[str]:
-        """Return the list of supported features."""
-        # TODO:
-        return set()
-
     @property
     def has_emeter(self) -> bool:
         """Return if the device has emeter."""

kasa/tests/test_cli.py

@@ -37,6 +37,11 @@ async def test_update_called_by_cli(dev, mocker):
     """Test that device update is called on main."""
     runner = CliRunner()
     update = mocker.patch.object(dev, "update")
+
+    # These will mock the features to avoid accessing non-existing
+    mocker.patch("kasa.device.Device.features", return_value={})
+    mocker.patch("kasa.iot.iotdevice.IotDevice.features", return_value={})
+
     mocker.patch("kasa.discover.Discover.discover_single", return_value=dev)
 
     res = await runner.invoke(
@@ -49,6 +54,7 @@ async def test_update_called_by_cli(dev, mocker):
             "--password",
             "bar",
         ],
+        catch_exceptions=False,
     )
     assert res.exit_code == 0
     update.assert_called()
@@ -292,6 +298,10 @@ async def test_brightness(dev):
 async def test_json_output(dev: Device, mocker):
     """Test that the json output produces correct output."""
     mocker.patch("kasa.Discover.discover", return_value={"127.0.0.1": dev})
+    # These will mock the features to avoid accessing non-existing
+    mocker.patch("kasa.device.Device.features", return_value={})
+    mocker.patch("kasa.iot.iotdevice.IotDevice.features", return_value={})
+
     runner = CliRunner()
     res = await runner.invoke(cli, ["--json", "state"], obj=dev)
     assert res.exit_code == 0
@@ -345,6 +355,10 @@ async def test_without_device_type(dev, mocker):
     discovery_mock = mocker.patch(
         "kasa.discover.Discover.discover_single", return_value=dev
     )
+    # These will mock the features to avoid accessing non-existing
+    mocker.patch("kasa.device.Device.features", return_value={})
+    mocker.patch("kasa.iot.iotdevice.IotDevice.features", return_value={})
+
     res = await runner.invoke(
         cli,
         [
@@ -410,6 +424,10 @@ async def test_duplicate_target_device():
 
 async def test_discover(discovery_mock, mocker):
     """Test discovery output."""
+    # These will mock the features to avoid accessing non-existing
+    mocker.patch("kasa.device.Device.features", return_value={})
+    mocker.patch("kasa.iot.iotdevice.IotDevice.features", return_value={})
+
     runner = CliRunner()
     res = await runner.invoke(
         cli,
@@ -429,6 +447,10 @@ async def test_discover(discovery_mock, mocker):
 
 async def test_discover_host(discovery_mock, mocker):
     """Test discovery output."""
+    # These will mock the features to avoid accessing non-existing
+    mocker.patch("kasa.device.Device.features", return_value={})
+    mocker.patch("kasa.iot.iotdevice.IotDevice.features", return_value={})
+
     runner = CliRunner()
     res = await runner.invoke(
         cli,

kasa/tests/test_feature.py

@@ -0,0 +1,79 @@
+import pytest
+
+from kasa import Feature, FeatureType
+
+
+@pytest.fixture
+def dummy_feature() -> Feature:
+    # create_autospec for device slows tests way too much, so we use a dummy here
+    class DummyDevice:
+        pass
+
+    feat = Feature(
+        device=DummyDevice(),  # type: ignore[arg-type]
+        name="dummy_feature",
+        attribute_getter="dummygetter",
+        attribute_setter="dummysetter",
+        container=None,
+        icon="mdi:dummy",
+        type=FeatureType.BinarySensor,
+    )
+    return feat
+
+
+def test_feature_api(dummy_feature: Feature):
+    """Test all properties of a dummy feature."""
+    assert dummy_feature.device is not None
+    assert dummy_feature.name == "dummy_feature"
+    assert dummy_feature.attribute_getter == "dummygetter"
+    assert dummy_feature.attribute_setter == "dummysetter"
+    assert dummy_feature.container is None
+    assert dummy_feature.icon == "mdi:dummy"
+    assert dummy_feature.type == FeatureType.BinarySensor
+
+
+def test_feature_value(dummy_feature: Feature):
+    """Verify that property gets accessed on *value* access."""
+    dummy_feature.attribute_getter = "test_prop"
+    dummy_feature.device.test_prop = "dummy"  # type: ignore[attr-defined]
+    assert dummy_feature.value == "dummy"
+
+
+def test_feature_value_container(mocker, dummy_feature: Feature):
+    """Test that container's attribute is accessed when expected."""
+
+    class DummyContainer:
+        @property
+        def test_prop(self):
+            return "dummy"
+
+    dummy_feature.container = DummyContainer()
+    dummy_feature.attribute_getter = "test_prop"
+
+    mock_dev_prop = mocker.patch.object(
+        dummy_feature, "test_prop", new_callable=mocker.PropertyMock, create=True
+    )
+
+    assert dummy_feature.value == "dummy"
+    mock_dev_prop.assert_not_called()
+
+
+def test_feature_value_callable(dev, dummy_feature: Feature):
+    """Verify that callables work as *attribute_getter*."""
+    dummy_feature.attribute_getter = lambda x: "dummy value"
+    assert dummy_feature.value == "dummy value"
+
+
+async def test_feature_setter(dev, mocker, dummy_feature: Feature):
+    """Verify that *set_value* calls the defined method."""
+    mock_set_dummy = mocker.patch.object(dummy_feature.device, "set_dummy", create=True)
+    dummy_feature.attribute_setter = "set_dummy"
+    await dummy_feature.set_value("dummy value")
+    mock_set_dummy.assert_called_with("dummy value")
+
+
+async def test_feature_setter_read_only(dummy_feature):
+    """Verify that read-only feature raises an exception when trying to change it."""
+    dummy_feature.attribute_setter = None
+    with pytest.raises(ValueError):
+        await dummy_feature.set_value("value for read only feature")

kasa/tests/test_smartdevice.py

@@ -67,7 +67,7 @@ async def test_invalid_connection(dev):
 async def test_initial_update_emeter(dev, mocker):
     """Test that the initial update performs second query if emeter is available."""
     dev._last_update = None
-    dev._features = set()
+    dev._legacy_features = set()
     spy = mocker.spy(dev.protocol, "query")
     await dev.update()
     # Devices with small buffers may require 3 queries
@@ -79,7 +79,7 @@ async def test_initial_update_emeter(dev, mocker):
 async def test_initial_update_no_emeter(dev, mocker):
     """Test that the initial update performs second query if emeter is available."""
     dev._last_update = None
-    dev._features = set()
+    dev._legacy_features = set()
     spy = mocker.spy(dev.protocol, "query")
     await dev.update()
     # 2 calls are necessary as some devices crash on unexpected modules
@@ -218,9 +218,9 @@ async def test_features(dev):
     """Make sure features is always accessible."""
     sysinfo = dev._last_update["system"]["get_sysinfo"]
     if "feature" in sysinfo:
-        assert dev.features == set(sysinfo["feature"].split(":"))
+        assert dev._legacy_features == set(sysinfo["feature"].split(":"))
     else:
-        assert dev.features == set()
+        assert dev._legacy_features == set()
 
 
 @device_iot