python-kasa · sdb9696 · May 10, 2024 · May 9, 2024 · May 10, 2024 · May 10, 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
@@ -40,6 +40,7 @@
     IotProtocol,
     _deprecated_TPLinkSmartHomeProtocol,  # noqa: F401
 )
+from kasa.module import Module
 from kasa.protocol import BaseProtocol
 from kasa.smartprotocol import SmartProtocol
 
@@ -60,6 +61,7 @@
     "Device",
     "Bulb",
     "Plug",
+    "Module",
     "KasaException",
     "AuthenticationError",
     "DeviceError",

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/modules/ledmodule.py → kasa/interfaces/ledinterface.py b/kasa/modules/ledmodule.py → kasa/interfaces/ledinterface.py
@@ -2,17 +2,15 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+
 from ..feature import Feature
 from ..module import Module
 
 
-class LedModule(Module):
+class LedInterface(Module, ABC):
     """Base interface to represent a LED module."""
 
-    # This needs to implement abstract methods for typing to work with
-    # overload get_module(type[ModuleT]) -> ModuleT:
-    # https://discuss.python.org/t/add-abstracttype-to-the-typing-module/21996
-
     def _initialize_features(self):
         """Initialize features."""
         device = self._device
@@ -22,7 +20,7 @@ def _initialize_features(self):
                 container=self,
                 name="LED",
                 id="led",
-                icon="mdi:led-{state}",
+                icon="mdi:led",
                 attribute_getter="led",
                 attribute_setter="set_led",
                 type=Feature.Type.Switch,
@@ -31,19 +29,10 @@ def _initialize_features(self):
         )
 
     @property
+    @abstractmethod
     def led(self) -> bool:
         """Return current led status."""
-        raise NotImplementedError()
 
+    @abstractmethod
     async def set_led(self, enable: bool) -> None:
         """Set led."""
-        raise NotImplementedError()
-
-    def query(self) -> dict:
-        """Query to execute during the update cycle."""
-        raise NotImplementedError()
-
-    @property
-    def data(self):
-        """Query to execute during the update cycle."""
-        raise NotImplementedError()
diff --git a/kasa/modules/lighteffectmodule.py → kasa/interfaces/lighteffectinterface.py b/kasa/modules/lighteffectmodule.py → kasa/interfaces/lighteffectinterface.py
@@ -2,17 +2,15 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+
 from ..feature import Feature
 from ..module import Module
 
 
-class LightEffectModule(Module):
+class LightEffectInterface(Module, ABC):
     """Interface to represent a light effect module."""
 
-    # This needs to implement abstract methods for typing to work with
-    # overload get_module(type[ModuleT]) -> ModuleT:
-    # https://discuss.python.org/t/add-abstracttype-to-the-typing-module/21996
-
     LIGHT_EFFECTS_OFF = "Off"
 
     def _initialize_features(self):
@@ -33,24 +31,25 @@ def _initialize_features(self):
         )
 
     @property
+    @abstractmethod
     def has_custom_effects(self) -> bool:
         """Return True if the device supports setting custom effects."""
-        raise NotImplementedError()
 
     @property
+    @abstractmethod
     def effect(self) -> str:
         """Return effect state or name."""
-        raise NotImplementedError()
 
     @property
+    @abstractmethod
     def effect_list(self) -> list[str]:
         """Return built-in effects list.
 
         Example:
             ['Aurora', 'Bubbling Cauldron', ...]
         """
-        raise NotImplementedError()
 
+    @abstractmethod
     async def set_effect(
         self,
         effect: str,
@@ -70,7 +69,6 @@ async def set_effect(
         :param int brightness: The wanted brightness
         :param int transition: The wanted transition time
         """
-        raise NotImplementedError()
 
     async def set_custom_effect(
         self,
@@ -80,13 +78,3 @@ async def set_custom_effect(
 
         :param str effect_dict: The custom effect dict to set
         """
-        raise NotImplementedError()
-
-    def query(self) -> dict:
-        """Query to execute during the update cycle."""
-        raise NotImplementedError()
-
-    @property
-    def data(self):
-        """Query to execute during the update cycle."""
-        raise NotImplementedError()
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 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 Module, ModuleT
+from ..module import Module
+from ..modulemapping import ModuleMapping, ModuleName
 from ..protocol import BaseProtocol
 from .iotmodule import IotModule
 from .modules import Emeter, Time
@@ -190,44 +191,26 @@ 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."""
-        return self._modules
+        return cast(ModuleMapping[IotModule], 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 | Module | 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, Module):
-            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
@@ -289,11 +272,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
@@ -368,15 +351,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
 
             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

diff --git a/kasa/iot/iotlightstrip.py b/kasa/iot/iotlightstrip.py
@@ -4,6 +4,7 @@
 
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
+from ..module import Module
 from ..protocol import BaseProtocol
 from .effects import EFFECT_NAMES_V1
 from .iotbulb import IotBulb
@@ -55,10 +56,10 @@ def __init__(
     ) -> None:
         super().__init__(host=host, config=config, protocol=protocol)
         self._device_type = DeviceType.LightStrip
-        self._light_effect_module = LightEffectModule(
-            self, "smartlife.iot.lighting_effect"
+        self.add_module(
+            Module.LightEffect,
+            LightEffectModule(self, "smartlife.iot.lighting_effect"),
         )
-        self.add_module("lighteffectmodule", self._light_effect_module)
 
     @property  # type: ignore
     @requires_update
@@ -79,7 +80,7 @@ def effect(self) -> dict:
              'name': ''}
         """
         # LightEffectModule returns the current effect name
-        # so return the dict here for backwards compatability
+        # so return the dict here for backwards compatibility
         return self.sys_info["lighting_effect_state"]
 
     @property  # type: ignore
@@ -91,7 +92,7 @@ def effect_list(self) -> list[str] | None:
             ['Aurora', 'Bubbling Cauldron', ...]
         """
         # LightEffectModule returns effect names along with a LIGHT_EFFECTS_OFF value
-        # so return the original effect names here for backwards compatability
+        # so return the original effect names here for backwards compatibility
         return EFFECT_NAMES_V1 if self.has_effects else None
 
     @requires_update
@@ -114,7 +115,7 @@ async def set_effect(
         :param int brightness: The wanted brightness
         :param int transition: The wanted transition time
         """
-        await self._light_effect_module.set_effect(
+        await self.modules[Module.LightEffect].set_effect(
             effect, brightness=brightness, transition=transition
         )
 
@@ -129,4 +130,4 @@ async def set_custom_effect(
         """
         if not self.has_effects:
             raise KasaException("Bulb does not support effects.")
-        await self._light_effect_module.set_custom_effect(effect_dict)
+        await self.modules[Module.LightEffect].set_custom_effect(effect_dict)
diff --git a/kasa/iot/iotplug.py b/kasa/iot/iotplug.py
@@ -6,6 +6,7 @@
 
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
+from ..module import Module
 from ..protocol import BaseProtocol
 from .iotdevice import IotDevice, requires_update
 from .modules import Antitheft, Cloud, LedModule, Schedule, Time, Usage
@@ -57,8 +58,7 @@ def __init__(
         self.add_module("antitheft", Antitheft(self, "anti_theft"))
         self.add_module("time", Time(self, "time"))
         self.add_module("cloud", Cloud(self, "cnCloud"))
-        self._led_module = LedModule(self, "system")
-        self.add_module("ledmodule", self._led_module)
+        self.add_module(Module.Led, LedModule(self, "system"))
 
     @property  # type: ignore
     @requires_update
@@ -79,11 +79,11 @@ async def turn_off(self, **kwargs):
     @requires_update
     def led(self) -> bool:
         """Return the state of the led."""
-        return self._led_module.led
+        return self.modules[Module.Led].led
 
     async def set_led(self, state: bool):
         """Set the state of the led (night mode)."""
-        return await self._led_module.set_led(state)
+        return await self.modules[Module.Led].set_led(state)
 
 
 class IotWallSwitch(IotPlug):