Create common interfaces for remaining device types (#895)

Introduce common module interfaces across smart and iot devices and provide better typing implementation for getting modules to support this.

kasa/iot/effects.py

@@ -0,0 +1,298 @@
+"""Module for light strip effects (LB*, KL*, KB*)."""
+
+from __future__ import annotations
+
+from typing import cast
+
+EFFECT_AURORA = {
+    "custom": 0,
+    "id": "xqUxDhbAhNLqulcuRMyPBmVGyTOyEMEu",
+    "brightness": 100,
+    "name": "Aurora",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "sequence",
+    "duration": 0,
+    "transition": 1500,
+    "direction": 4,
+    "spread": 7,
+    "repeat_times": 0,
+    "sequence": [[120, 100, 100], [240, 100, 100], [260, 100, 100], [280, 100, 100]],
+}
+EFFECT_BUBBLING_CAULDRON = {
+    "custom": 0,
+    "id": "tIwTRQBqJpeNKbrtBMFCgkdPTbAQGfRP",
+    "brightness": 100,
+    "name": "Bubbling Cauldron",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [100, 270],
+    "saturation_range": [80, 100],
+    "brightness_range": [50, 100],
+    "duration": 0,
+    "transition": 200,
+    "init_states": [[270, 100, 100]],
+    "fadeoff": 1000,
+    "random_seed": 24,
+    "backgrounds": [[270, 40, 50]],
+}
+EFFECT_CANDY_CANE = {
+    "custom": 0,
+    "id": "HCOttllMkNffeHjEOLEgrFJjbzQHoxEJ",
+    "brightness": 100,
+    "name": "Candy Cane",
+    "segments": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "sequence",
+    "duration": 700,
+    "transition": 500,
+    "direction": 1,
+    "spread": 1,
+    "repeat_times": 0,
+    "sequence": [
+        [0, 0, 100],
+        [0, 0, 100],
+        [360, 81, 100],
+        [0, 0, 100],
+        [0, 0, 100],
+        [360, 81, 100],
+        [360, 81, 100],
+        [0, 0, 100],
+        [0, 0, 100],
+        [360, 81, 100],
+        [360, 81, 100],
+        [360, 81, 100],
+        [360, 81, 100],
+        [0, 0, 100],
+        [0, 0, 100],
+        [360, 81, 100],
+    ],
+}
+EFFECT_CHRISTMAS = {
+    "custom": 0,
+    "id": "bwTatyinOUajKrDwzMmqxxJdnInQUgvM",
+    "brightness": 100,
+    "name": "Christmas",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [136, 146],
+    "saturation_range": [90, 100],
+    "brightness_range": [50, 100],
+    "duration": 5000,
+    "transition": 0,
+    "init_states": [[136, 0, 100]],
+    "fadeoff": 2000,
+    "random_seed": 100,
+    "backgrounds": [[136, 98, 75], [136, 0, 0], [350, 0, 100], [350, 97, 94]],
+}
+EFFECT_FLICKER = {
+    "custom": 0,
+    "id": "bCTItKETDFfrKANolgldxfgOakaarARs",
+    "brightness": 100,
+    "name": "Flicker",
+    "segments": [1],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [30, 40],
+    "saturation_range": [100, 100],
+    "brightness_range": [50, 100],
+    "duration": 0,
+    "transition": 0,
+    "transition_range": [375, 500],
+    "init_states": [[30, 81, 80]],
+}
+EFFECT_HANUKKAH = {
+    "custom": 0,
+    "id": "CdLeIgiKcQrLKMINRPTMbylATulQewLD",
+    "brightness": 100,
+    "name": "Hanukkah",
+    "segments": [1],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [200, 210],
+    "saturation_range": [0, 100],
+    "brightness_range": [50, 100],
+    "duration": 1500,
+    "transition": 0,
+    "transition_range": [400, 500],
+    "init_states": [[35, 81, 80]],
+}
+EFFECT_HAUNTED_MANSION = {
+    "custom": 0,
+    "id": "oJnFHsVQzFUTeIOBAhMRfVeujmSauhjJ",
+    "brightness": 80,
+    "name": "Haunted Mansion",
+    "segments": [80],
+    "expansion_strategy": 2,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [45, 45],
+    "saturation_range": [10, 10],
+    "brightness_range": [0, 80],
+    "duration": 0,
+    "transition": 0,
+    "transition_range": [50, 1500],
+    "init_states": [[45, 10, 100]],
+    "fadeoff": 200,
+    "random_seed": 1,
+    "backgrounds": [[45, 10, 100]],
+}
+EFFECT_ICICLE = {
+    "custom": 0,
+    "id": "joqVjlaTsgzmuQQBAlHRkkPAqkBUiqeb",
+    "brightness": 70,
+    "name": "Icicle",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "sequence",
+    "duration": 0,
+    "transition": 400,
+    "direction": 4,
+    "spread": 3,
+    "repeat_times": 0,
+    "sequence": [
+        [190, 100, 70],
+        [190, 100, 70],
+        [190, 30, 50],
+        [190, 100, 70],
+        [190, 100, 70],
+    ],
+}
+EFFECT_LIGHTNING = {
+    "custom": 0,
+    "id": "ojqpUUxdGHoIugGPknrUcRoyJiItsjuE",
+    "brightness": 100,
+    "name": "Lightning",
+    "segments": [7, 20, 23, 32, 34, 35, 49, 65, 66, 74, 80],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [240, 240],
+    "saturation_range": [10, 11],
+    "brightness_range": [90, 100],
+    "duration": 0,
+    "transition": 50,
+    "init_states": [[240, 30, 100]],
+    "fadeoff": 150,
+    "random_seed": 600,
+    "backgrounds": [[200, 100, 100], [200, 50, 10], [210, 10, 50], [240, 10, 0]],
+}
+EFFECT_OCEAN = {
+    "custom": 0,
+    "id": "oJjUMosgEMrdumfPANKbkFmBcAdEQsPy",
+    "brightness": 30,
+    "name": "Ocean",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "sequence",
+    "duration": 0,
+    "transition": 2000,
+    "direction": 3,
+    "spread": 16,
+    "repeat_times": 0,
+    "sequence": [[198, 84, 30], [198, 70, 30], [198, 10, 30]],
+}
+EFFECT_RAINBOW = {
+    "custom": 0,
+    "id": "izRhLCQNcDzIKdpMPqSTtBMuAIoreAuT",
+    "brightness": 100,
+    "name": "Rainbow",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "sequence",
+    "duration": 0,
+    "transition": 1500,
+    "direction": 1,
+    "spread": 12,
+    "repeat_times": 0,
+    "sequence": [[0, 100, 100], [100, 100, 100], [200, 100, 100], [300, 100, 100]],
+}
+EFFECT_RAINDROP = {
+    "custom": 0,
+    "id": "QbDFwiSFmLzQenUOPnJrsGqyIVrJrRsl",
+    "brightness": 30,
+    "name": "Raindrop",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [200, 200],
+    "saturation_range": [10, 20],
+    "brightness_range": [10, 30],
+    "duration": 0,
+    "transition": 1000,
+    "init_states": [[200, 40, 100]],
+    "fadeoff": 1000,
+    "random_seed": 24,
+    "backgrounds": [[200, 40, 0]],
+}
+EFFECT_SPRING = {
+    "custom": 0,
+    "id": "URdUpEdQbnOOechDBPMkKrwhSupLyvAg",
+    "brightness": 100,
+    "name": "Spring",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [0, 90],
+    "saturation_range": [30, 100],
+    "brightness_range": [90, 100],
+    "duration": 600,
+    "transition": 0,
+    "transition_range": [2000, 6000],
+    "init_states": [[80, 30, 100]],
+    "fadeoff": 1000,
+    "random_seed": 20,
+    "backgrounds": [[130, 100, 40]],
+}
+EFFECT_VALENTINES = {
+    "custom": 0,
+    "id": "QglBhMShPHUAuxLqzNEefFrGiJwahOmz",
+    "brightness": 100,
+    "name": "Valentines",
+    "segments": [0],
+    "expansion_strategy": 1,
+    "enable": 1,
+    "type": "random",
+    "hue_range": [340, 340],
+    "saturation_range": [30, 40],
+    "brightness_range": [90, 100],
+    "duration": 600,
+    "transition": 2000,
+    "init_states": [[340, 30, 100]],
+    "fadeoff": 3000,
+    "random_seed": 100,
+    "backgrounds": [[340, 20, 50], [20, 50, 50], [0, 100, 50]],
+}
+
+EFFECTS_LIST_V1 = [
+    EFFECT_AURORA,
+    EFFECT_BUBBLING_CAULDRON,
+    EFFECT_CANDY_CANE,
+    EFFECT_CHRISTMAS,
+    EFFECT_FLICKER,
+    EFFECT_HANUKKAH,
+    EFFECT_HAUNTED_MANSION,
+    EFFECT_ICICLE,
+    EFFECT_LIGHTNING,
+    EFFECT_OCEAN,
+    EFFECT_RAINBOW,
+    EFFECT_RAINDROP,
+    EFFECT_SPRING,
+    EFFECT_VALENTINES,
+]
+
+EFFECT_NAMES_V1: list[str] = [cast(str, effect["name"]) for effect in EFFECTS_LIST_V1]
+EFFECT_MAPPING_V1 = {effect["name"]: effect for effect in EFFECTS_LIST_V1}

kasa/iot/iotdevice.py

@@ -19,14 +19,15 @@ import functools
 import inspect
 import logging
 from datetime import datetime, timedelta
-from typing import Any, Mapping, Sequence, cast, overload
+from typing import TYPE_CHECKING, Any, Mapping, Sequence, cast
 
 from ..device import Device, WifiNetwork
 from ..deviceconfig import DeviceConfig
 from ..emeterstatus import EmeterStatus
 from ..exceptions import KasaException
 from ..feature import Feature
-from ..module import ModuleT
+from ..module import Module
+from ..modulemapping import ModuleMapping, ModuleName
 from ..protocol import BaseProtocol
 from .iotmodule import IotModule
 from .modules import Emeter, Time
@@ -190,7 +191,7 @@ class IotDevice(Device):
         self._supported_modules: dict[str, IotModule] | None = None
         self._legacy_features: set[str] = set()
         self._children: Mapping[str, IotDevice] = {}
-        self._modules: dict[str, IotModule] = {}
+        self._modules: dict[str | ModuleName[Module], IotModule] = {}
 
     @property
     def children(self) -> Sequence[IotDevice]:
@@ -198,38 +199,20 @@ class IotDevice(Device):
         return list(self._children.values())
 
     @property
-    def modules(self) -> dict[str, IotModule]:
+    def modules(self) -> ModuleMapping[IotModule]:
         """Return the device modules."""
+        if TYPE_CHECKING:
+            return cast(ModuleMapping[IotModule], self._modules)
         return self._modules
 
-    @overload
-    def get_module(self, module_type: type[ModuleT]) -> ModuleT | None: ...
-
-    @overload
-    def get_module(self, module_type: str) -> IotModule | None: ...
-
-    def get_module(
-        self, module_type: type[ModuleT] | str
-    ) -> ModuleT | IotModule | None:
-        """Return the module from the device modules or None if not present."""
-        if isinstance(module_type, str):
-            module_name = module_type.lower()
-        elif issubclass(module_type, IotModule):
-            module_name = module_type.__name__.lower()
-        else:
-            return None
-        if module_name in self.modules:
-            return self.modules[module_name]
-        return None
-
-    def add_module(self, name: str, module: IotModule):
+    def add_module(self, name: str | ModuleName[Module], module: IotModule):
         """Register a module."""
         if name in self.modules:
             _LOGGER.debug("Module %s already registered, ignoring..." % name)
             return
 
         _LOGGER.debug("Adding module %s", module)
-        self.modules[name] = module
+        self._modules[name] = module
 
     def _create_request(
         self, target: str, cmd: str, arg: dict | None = None, child_ids=None
@@ -291,11 +274,11 @@ class IotDevice(Device):
 
     @property  # type: ignore
     @requires_update
-    def supported_modules(self) -> list[str]:
+    def supported_modules(self) -> list[str | ModuleName[Module]]:
         """Return a set of modules supported by the device."""
         # TODO: this should rather be called `features`, but we don't want to break
         #       the API now. Maybe just deprecate it and point the users to use this?
-        return list(self.modules.keys())
+        return list(self._modules.keys())
 
     @property  # type: ignore
     @requires_update
@@ -324,10 +307,11 @@ class IotDevice(Device):
             self._last_update = response
             self._set_sys_info(response["system"]["get_sysinfo"])
 
+        await self._modular_update(req)
+
         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):
@@ -352,6 +336,11 @@ class IotDevice(Device):
                 )
             )
 
+        for module in self._modules.values():
+            module._initialize_features()
+            for module_feat in module._module_features.values():
+                self._add_feature(module_feat)
+
     async def _modular_update(self, req: dict) -> None:
         """Execute an update query."""
         if self.has_emeter:
@@ -364,17 +353,15 @@ class IotDevice(Device):
         #  making separate handling for this unnecessary
         if self._supported_modules is None:
             supported = {}
-            for module in self.modules.values():
+            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():
+        for module in self._modules.values():
             if not module.is_supported:
                 _LOGGER.debug("Module %s not supported, skipping" % module)
                 continue

kasa/iot/iotlightstrip.py

@@ -4,10 +4,12 @@ from __future__ import annotations
 
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
-from ..effects import EFFECT_MAPPING_V1, EFFECT_NAMES_V1
+from ..module import Module
 from ..protocol import BaseProtocol
+from .effects import EFFECT_NAMES_V1
 from .iotbulb import IotBulb
 from .iotdevice import KasaException, requires_update
+from .modules.lighteffectmodule import LightEffectModule
 
 
 class IotLightStrip(IotBulb):
@@ -54,6 +56,10 @@ class IotLightStrip(IotBulb):
     ) -> None:
         super().__init__(host=host, config=config, protocol=protocol)
         self._device_type = DeviceType.LightStrip
+        self.add_module(
+            Module.LightEffect,
+            LightEffectModule(self, "smartlife.iot.lighting_effect"),
+        )
 
     @property  # type: ignore
     @requires_update
@@ -73,6 +79,8 @@ class IotLightStrip(IotBulb):
              'id': '',
              'name': ''}
         """
+        # LightEffectModule returns the current effect name
+        # so return the dict here for backwards compatibility
         return self.sys_info["lighting_effect_state"]
 
     @property  # type: ignore
@@ -83,6 +91,8 @@ class IotLightStrip(IotBulb):
         Example:
             ['Aurora', 'Bubbling Cauldron', ...]
         """
+        # LightEffectModule returns effect names along with a LIGHT_EFFECTS_OFF value
+        # so return the original effect names here for backwards compatibility
         return EFFECT_NAMES_V1 if self.has_effects else None
 
     @requires_update
@@ -105,15 +115,9 @@ class IotLightStrip(IotBulb):
         :param int brightness: The wanted brightness
         :param int transition: The wanted transition time
         """
-        if effect not in EFFECT_MAPPING_V1:
-            raise KasaException(f"The effect {effect} is not a built in effect.")
-        effect_dict = EFFECT_MAPPING_V1[effect]
-        if brightness is not None:
-            effect_dict["brightness"] = brightness
-        if transition is not None:
-            effect_dict["transition"] = transition
-
-        await self.set_custom_effect(effect_dict)
+        await self.modules[Module.LightEffect].set_effect(
+            effect, brightness=brightness, transition=transition
+        )
 
     @requires_update
     async def set_custom_effect(
@@ -126,8 +130,4 @@ class IotLightStrip(IotBulb):
         """
         if not self.has_effects:
             raise KasaException("Bulb does not support effects.")
-        await self._query_helper(
-            "smartlife.iot.lighting_effect",
-            "set_lighting_effect",
-            effect_dict,
-        )
+        await self.modules[Module.LightEffect].set_custom_effect(effect_dict)

kasa/iot/iotmodule.py

@@ -43,13 +43,19 @@ class IotModule(Module):
     @property
     def data(self):
         """Return the module specific raw data from the last update."""
-        if self._module not in self._device._last_update:
+        dev = self._device
+        q = self.query()
+
+        if not q:
+            return dev.sys_info
+
+        if self._module not in dev._last_update:
             raise KasaException(
                 f"You need to call update() prior accessing module data"
                 f" for '{self._module}'"
             )
 
-        return self._device._last_update[self._module]
+        return dev._last_update[self._module]
 
     @property
     def is_supported(self) -> bool:

kasa/iot/iotplug.py

@@ -6,10 +6,10 @@ import logging
 
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
-from ..feature import Feature
+from ..module import Module
 from ..protocol import BaseProtocol
 from .iotdevice import IotDevice, requires_update
-from .modules import Antitheft, Cloud, Schedule, Time, Usage
+from .modules import Antitheft, Cloud, LedModule, Schedule, Time, Usage
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -58,21 +58,7 @@ class IotPlug(IotDevice):
         self.add_module("antitheft", Antitheft(self, "anti_theft"))
         self.add_module("time", Time(self, "time"))
         self.add_module("cloud", Cloud(self, "cnCloud"))
-
-    async def _initialize_features(self):
-        await super()._initialize_features()
-
-        self._add_feature(
-            Feature(
-                device=self,
-                id="led",
-                name="LED",
-                icon="mdi:led-{state}",
-                attribute_getter="led",
-                attribute_setter="set_led",
-                type=Feature.Type.Switch,
-            )
-        )
+        self.add_module(Module.Led, LedModule(self, "system"))
 
     @property  # type: ignore
     @requires_update
@@ -93,14 +79,11 @@ class IotPlug(IotDevice):
     @requires_update
     def led(self) -> bool:
         """Return the state of the led."""
-        sys_info = self.sys_info
-        return bool(1 - sys_info["led_off"])
+        return self.modules[Module.Led].led
 
     async def set_led(self, state: bool):
         """Set the state of the led (night mode)."""
-        return await self._query_helper(
-            "system", "set_led_off", {"off": int(not state)}
-        )
+        return await self.modules[Module.Led].set_led(state)
 
 
 class IotWallSwitch(IotPlug):

kasa/iot/iotstrip.py

@@ -253,7 +253,6 @@ class IotStripPlug(IotPlug):
         self._last_update = parent._last_update
         self._set_sys_info(parent.sys_info)
         self._device_type = DeviceType.StripSocket
-        self._modules = {}
         self.protocol = parent.protocol  # Must use the same connection as the parent
         self.add_module("time", Time(self, "time"))
 

kasa/iot/modules/__init__.py

@@ -5,6 +5,7 @@ from .antitheft import Antitheft
 from .cloud import Cloud
 from .countdown import Countdown
 from .emeter import Emeter
+from .ledmodule import LedModule
 from .motion import Motion
 from .rulemodule import Rule, RuleModule
 from .schedule import Schedule
@@ -17,6 +18,7 @@ __all__ = [
     "Cloud",
     "Countdown",
     "Emeter",
+    "LedModule",
     "Motion",
     "Rule",
     "RuleModule",

kasa/iot/modules/ledmodule.py

@@ -0,0 +1,32 @@
+"""Module for led controls."""
+
+from __future__ import annotations
+
+from ...interfaces.led import Led
+from ..iotmodule import IotModule
+
+
+class LedModule(IotModule, Led):
+    """Implementation of led controls."""
+
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
+    @property
+    def mode(self):
+        """LED mode setting.
+
+        "always", "never"
+        """
+        return "always" if self.led else "never"
+
+    @property
+    def led(self) -> bool:
+        """Return the state of the led."""
+        sys_info = self.data
+        return bool(1 - sys_info["led_off"])
+
+    async def set_led(self, state: bool):
+        """Set the state of the led (night mode)."""
+        return await self.call("set_led_off", {"off": int(not state)})

kasa/iot/modules/lighteffectmodule.py

@@ -0,0 +1,97 @@
+"""Module for light effects."""
+
+from __future__ import annotations
+
+from ...interfaces.lighteffect import LightEffect
+from ..effects import EFFECT_MAPPING_V1, EFFECT_NAMES_V1
+from ..iotmodule import IotModule
+
+
+class LightEffectModule(IotModule, LightEffect):
+    """Implementation of dynamic light effects."""
+
+    @property
+    def effect(self) -> str:
+        """Return effect state.
+
+        Example:
+            {'brightness': 50,
+             'custom': 0,
+             'enable': 0,
+             'id': '',
+             'name': ''}
+        """
+        if (
+            (state := self.data.get("lighting_effect_state"))
+            and state.get("enable")
+            and (name := state.get("name"))
+            and name in EFFECT_NAMES_V1
+        ):
+            return name
+        return self.LIGHT_EFFECTS_OFF
+
+    @property
+    def effect_list(self) -> list[str]:
+        """Return built-in effects list.
+
+        Example:
+            ['Aurora', 'Bubbling Cauldron', ...]
+        """
+        effect_list = [self.LIGHT_EFFECTS_OFF]
+        effect_list.extend(EFFECT_NAMES_V1)
+        return effect_list
+
+    async def set_effect(
+        self,
+        effect: str,
+        *,
+        brightness: int | None = None,
+        transition: int | None = None,
+    ) -> None:
+        """Set an effect on the device.
+
+        If brightness or transition is defined,
+        its value will be used instead of the effect-specific default.
+
+        See :meth:`effect_list` for available effects,
+        or use :meth:`set_custom_effect` for custom effects.
+
+        :param str effect: The effect to set
+        :param int brightness: The wanted brightness
+        :param int transition: The wanted transition time
+        """
+        if effect == self.LIGHT_EFFECTS_OFF:
+            effect_dict = dict(self.data["lighting_effect_state"])
+            effect_dict["enable"] = 0
+        elif effect not in EFFECT_MAPPING_V1:
+            raise ValueError(f"The effect {effect} is not a built in effect.")
+        else:
+            effect_dict = EFFECT_MAPPING_V1[effect]
+        if brightness is not None:
+            effect_dict["brightness"] = brightness
+        if transition is not None:
+            effect_dict["transition"] = transition
+
+        await self.set_custom_effect(effect_dict)
+
+    async def set_custom_effect(
+        self,
+        effect_dict: dict,
+    ) -> None:
+        """Set a custom effect on the device.
+
+        :param str effect_dict: The custom effect dict to set
+        """
+        return await self.call(
+            "set_lighting_effect",
+            effect_dict,
+        )
+
+    @property
+    def has_custom_effects(self) -> bool:
+        """Return True if the device supports setting custom effects."""
+        return True
+
+    def query(self):
+        """Return the base query."""
+        return {}