python-kasa · sdb9696 · May 10, 2024 · May 2, 2024 · May 3, 2024 · May 3, 2024
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -21,6 +21,11 @@ repos:
   hooks:
   - id: mypy
     additional_dependencies: [types-click]
+    exclude: |
+      (?x)^(
+          kasa/modulemapping\.py|
+      )$
+
 
 - repo: https://github.com/PyCQA/doc8
   rev: 'v1.1.1'

diff --git a/devtools/create_module_fixtures.py b/devtools/create_module_fixtures.py
@@ -19,7 +19,7 @@
 def create_fixtures(dev: IotDevice, outputdir: Path):
     """Iterate over supported modules and create version-specific fixture files."""
     for name, module in dev.modules.items():
-        module_dir = outputdir / name
+        module_dir = outputdir / str(name)
         if not module_dir.exists():
             module_dir.mkdir(exist_ok=True, parents=True)
 

diff --git a/kasa/__init__.py b/kasa/__init__.py
@@ -16,7 +16,7 @@
 from typing import TYPE_CHECKING
 from warnings import warn
 
-from kasa.bulb import Bulb
+from kasa.bulb import Bulb, BulbPreset
 from kasa.credentials import Credentials
 from kasa.device import Device
 from kasa.device_type import DeviceType
@@ -36,12 +36,11 @@
     UnsupportedDeviceError,
 )
 from kasa.feature import Feature
-from kasa.iot.iotbulb import BulbPreset, TurnOnBehavior, TurnOnBehaviors
 from kasa.iotprotocol import (
     IotProtocol,
     _deprecated_TPLinkSmartHomeProtocol,  # noqa: F401
 )
-from kasa.plug import Plug
+from kasa.module import Module
 from kasa.protocol import BaseProtocol
 from kasa.smartprotocol import SmartProtocol
 
@@ -62,6 +61,7 @@
     "Device",
     "Bulb",
     "Plug",
+    "Module",
     "KasaException",
     "AuthenticationError",
     "DeviceError",

diff --git a/kasa/bulb.py b/kasa/bulb.py
@@ -54,11 +54,6 @@ def _raise_for_invalid_brightness(self, value):
     def is_color(self) -> bool:
         """Whether the bulb supports color changes."""
 
-    @property
-    @abstractmethod
-    def is_dimmable(self) -> bool:
-        """Whether the bulb supports brightness changes."""
-
     @property
     @abstractmethod
     def is_variable_color_temp(self) -> bool:

diff --git a/kasa/device.py b/kasa/device.py
@@ -6,7 +6,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Mapping, Sequence, overload
+from typing import TYPE_CHECKING, Any, Mapping, Sequence
 
 from .credentials import Credentials
 from .device_type import DeviceType
@@ -15,10 +15,13 @@
 from .exceptions import KasaException
 from .feature import Feature
 from .iotprotocol import IotProtocol
-from .module import Module, ModuleT
+from .module import Module
 from .protocol import BaseProtocol
 from .xortransport import XorTransport
 
+if TYPE_CHECKING:
+    from .modulemapping import ModuleMapping
+
 
 @dataclass
 class WifiNetwork:
@@ -113,21 +116,9 @@ async def disconnect(self):
 
     @property
     @abstractmethod
-    def modules(self) -> Mapping[str, Module]:
+    def modules(self) -> ModuleMapping[Module]:
         """Return the device modules."""
 
-    @overload
-    @abstractmethod
-    def get_module(self, module_type: type[ModuleT]) -> ModuleT | None: ...
-
-    @overload
-    @abstractmethod
-    def get_module(self, module_type: str) -> Module | None: ...
-
-    @abstractmethod
-    def get_module(self, module_type: type[ModuleT] | str) -> ModuleT | Module | None:
-        """Return the module from the device modules or None if not present."""
-
     @property
     @abstractmethod
     def is_on(self) -> bool:

diff --git a/kasa/interfaces/led.py b/kasa/interfaces/led.py
@@ -0,0 +1,38 @@
+"""Module for base light effect module."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from ..feature import Feature
+from ..module import Module
+
+
+class Led(Module, ABC):
+    """Base interface to represent a LED module."""
+
+    def _initialize_features(self):
+        """Initialize features."""
+        device = self._device
+        self._add_feature(
+            Feature(
+                device=device,
+                container=self,
+                name="LED",
+                id="led",
+                icon="mdi:led",
+                attribute_getter="led",
+                attribute_setter="set_led",
+                type=Feature.Type.Switch,
+                category=Feature.Category.Config,
+            )
+        )
+
+    @property
+    @abstractmethod
+    def led(self) -> bool:
+        """Return current led status."""
+
+    @abstractmethod
+    async def set_led(self, enable: bool) -> None:
+        """Set led."""
diff --git a/kasa/interfaces/lighteffect.py b/kasa/interfaces/lighteffect.py
@@ -0,0 +1,80 @@
+"""Module for base light effect module."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from ..feature import Feature
+from ..module import Module
+
+
+class LightEffect(Module, ABC):
+    """Interface to represent a light effect module."""
+
+    LIGHT_EFFECTS_OFF = "Off"
+
+    def _initialize_features(self):
+        """Initialize features."""
+        device = self._device
+        self._add_feature(
+            Feature(
+                device,
+                id="light_effect",
+                name="Light effect",
+                container=self,
+                attribute_getter="effect",
+                attribute_setter="set_effect",
+                category=Feature.Category.Primary,
+                type=Feature.Type.Choice,
+                choices_getter="effect_list",
+            )
+        )
+
+    @property
+    @abstractmethod
+    def has_custom_effects(self) -> bool:
+        """Return True if the device supports setting custom effects."""
+
+    @property
+    @abstractmethod
+    def effect(self) -> str:
+        """Return effect state or name."""
+
+    @property
+    @abstractmethod
+    def effect_list(self) -> list[str]:
+        """Return built-in effects list.
+
+        Example:
+            ['Aurora', 'Bubbling Cauldron', ...]
+        """
+
+    @abstractmethod
+    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
+        """
+
+    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
+        """
diff --git a/kasa/effects.py → kasa/iot/effects.py b/kasa/effects.py → kasa/iot/effects.py
diff --git a/kasa/iot/iotdevice.py b/kasa/iot/iotdevice.py
@@ -19,14 +19,15 @@
 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,46 +191,28 @@ def __init__(
         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]:
         """Return list of children."""
         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 @@ def features(self) -> dict[str, Feature]:
 
     @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 @@ async def update(self, update_children: bool = True):
             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 @@ async def _initialize_features(self):
                 )
             )
 
+        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 @@ async def _modular_update(self, req: dict) -> None:
         #  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