From f1e5b9653f6cbd2ff7ca2b74004b3bf3fe0f71b6 Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Tue, 20 Feb 2024 00:26:47 +0100 Subject: [PATCH 1/8] Generalize smartdevice child support * Initialize children's modules (and features) using the child component negotiation results * Set device_type based on the device response * Print out child features in cli 'state' * Add --child option to cli 'command' to allow targeting child devices * Guard "generic" features like rssi, ssid, etc. only to devices which have this information Note, we do not currently perform queries on child modules so some data may not be available. At the moment, a stop-gap solution to use parent's data is used but this is not always correct; even if the device shares the same clock and cloud connectivity, it may have its own firmware updates. --- kasa/cli.py | 20 ++++- kasa/device.py | 1 + kasa/device_type.py | 1 + kasa/smart/modules/__init__.py | 1 + kasa/smart/modules/childdevicemodule.py | 12 ++- kasa/smart/smartchilddevice.py | 37 ++++++-- kasa/smart/smartdevice.py | 110 +++++++++++++++--------- kasa/smart/smartmodule.py | 22 +++-- kasa/tests/test_childdevice.py | 1 - kasa/tests/test_smartdevice.py | 4 +- 10 files changed, 146 insertions(+), 63 deletions(-) diff --git a/kasa/cli.py b/kasa/cli.py index 4d3590d10..0a2fff8a1 100755 --- a/kasa/cli.py +++ b/kasa/cli.py @@ -558,9 +558,14 @@ async def state(ctx, dev: Device): echo(f"\tPort: {dev.port}") echo(f"\tDevice state: {dev.is_on}") if dev.is_strip: - echo("\t[bold]== Plugs ==[/bold]") - for plug in dev.children: # type: ignore - echo(f"\t* Socket '{plug.alias}' state: {plug.is_on} since {plug.on_since}") + echo("\t[bold]== Children ==[/bold]") + for child in dev.children: + echo(f"\t* {child.alias} ({child.model}, {child.device_type})") + for feat in child.features.values(): + try: + echo(f"\t\t{feat.name}: {feat.value}") + except Exception as ex: + echo(f"\t\t{feat.name}: got exception (%s)" % ex) echo() echo("\t[bold]== Generic information ==[/bold]") @@ -641,13 +646,20 @@ async def raw_command(ctx, dev: Device, module, command, parameters): @cli.command(name="command") @pass_dev @click.option("--module", required=False, help="Module for IOT protocol.") +@click.option("--child", required=False, help="Child ID for controlling sub-devices") @click.argument("command") @click.argument("parameters", default=None, required=False) -async def cmd_command(dev: Device, module, command, parameters): +async def cmd_command(dev: Device, module, child, command, parameters): """Run a raw command on the device.""" if parameters is not None: parameters = ast.literal_eval(parameters) + if child: + echo(f"Selecting child {child} from {dev}") + # TODO: Update required to initialize children + await dev.update() + dev = dev._children[child] + if isinstance(dev, IotDevice): res = await dev._query_helper(module, command, parameters) elif isinstance(dev, SmartDevice): diff --git a/kasa/device.py b/kasa/device.py index 3c38b5446..fb92f18e9 100644 --- a/kasa/device.py +++ b/kasa/device.py @@ -71,6 +71,7 @@ def __init__( self.modules: Dict[str, Any] = {} self._features: Dict[str, Feature] = {} + self._parent: Optional["Device"] = None @staticmethod async def connect( diff --git a/kasa/device_type.py b/kasa/device_type.py index 162fc4f27..41dd6e363 100755 --- a/kasa/device_type.py +++ b/kasa/device_type.py @@ -14,6 +14,7 @@ class DeviceType(Enum): StripSocket = "stripsocket" Dimmer = "dimmer" LightStrip = "lightstrip" + Sensor = "sensor" Unknown = "unknown" @staticmethod diff --git a/kasa/smart/modules/__init__.py b/kasa/smart/modules/__init__.py index 02c3b86af..c8e17c57f 100644 --- a/kasa/smart/modules/__init__.py +++ b/kasa/smart/modules/__init__.py @@ -8,6 +8,7 @@ from .ledmodule import LedModule from .lighttransitionmodule import LightTransitionModule from .timemodule import TimeModule +from .firmware import Firmware __all__ = [ "TimeModule", diff --git a/kasa/smart/modules/childdevicemodule.py b/kasa/smart/modules/childdevicemodule.py index 991acc25b..62e024d0c 100644 --- a/kasa/smart/modules/childdevicemodule.py +++ b/kasa/smart/modules/childdevicemodule.py @@ -1,4 +1,6 @@ """Implementation for child devices.""" +from typing import Dict + from ..smartmodule import SmartModule @@ -6,4 +8,12 @@ class ChildDeviceModule(SmartModule): """Implementation for child devices.""" REQUIRED_COMPONENT = "child_device" - QUERY_GETTER_NAME = "get_child_device_list" + + def query(self) -> Dict: + """Query to execute during the update cycle.""" + # TODO: There is no need to fetch the component list every time, + # so this should be optimized only for the init. + return { + "get_child_device_list": None, + "get_child_device_component_list": None, + } diff --git a/kasa/smart/smartchilddevice.py b/kasa/smart/smartchilddevice.py index 698982b67..6d7bfa587 100644 --- a/kasa/smart/smartchilddevice.py +++ b/kasa/smart/smartchilddevice.py @@ -1,4 +1,5 @@ """Child device implementation.""" +import logging from typing import Optional from ..device_type import DeviceType @@ -6,6 +7,8 @@ from ..smartprotocol import SmartProtocol, _ChildProtocolWrapper from .smartdevice import SmartDevice +_LOGGER = logging.getLogger(__name__) + class SmartChildDevice(SmartDevice): """Presentation of a child device. @@ -16,23 +19,41 @@ class SmartChildDevice(SmartDevice): def __init__( self, parent: SmartDevice, - child_id: str, + info, + component_info, config: Optional[DeviceConfig] = None, protocol: Optional[SmartProtocol] = None, ) -> None: super().__init__(parent.host, config=parent.config, protocol=parent.protocol) self._parent = parent - self._id = child_id - self.protocol = _ChildProtocolWrapper(child_id, parent.protocol) - self._device_type = DeviceType.StripSocket + self._update_internal_state(info) + self._components = component_info + self._id = info["device_id"] + self.protocol = _ChildProtocolWrapper(self._id, parent.protocol) async def update(self, update_children: bool = True): """Noop update. The parent updates our internals.""" - def update_internal_state(self, info): - """Set internal state for the child.""" - # TODO: cleanup the _last_update, _sys_info, _info, _data mess. - self._last_update = self._sys_info = self._info = info + @classmethod + async def create(cls, parent: SmartDevice, child_info, child_components): + """Create a child device based on device info and component listing.""" + child: "SmartChildDevice" = cls(parent, child_info, child_components) + await child._initialize_modules() + await child._initialize_features() + return child + + @property + def device_type(self) -> DeviceType: + """Return child device type.""" + child_device_map = { + "plug.powerstrip.sub-plug": DeviceType.Plug, + "subg.trigger.temp-hmdt-sensor": DeviceType.Sensor, + } + dev_type = child_device_map.get(self.sys_info["category"]) + if dev_type is None: + _LOGGER.warning("Unknown child device type, please open issue ") + dev_type = DeviceType.Unknown + return dev_type def __repr__(self): return f"" diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py index 62657d816..9e691c774 100644 --- a/kasa/smart/smartdevice.py +++ b/kasa/smart/smartdevice.py @@ -12,22 +12,13 @@ from ..exceptions import AuthenticationException, SmartDeviceException, SmartErrorCode from ..feature import Feature, FeatureType from ..smartprotocol import SmartProtocol -from .modules import ( # noqa: F401 - AutoOffModule, - ChildDeviceModule, - CloudModule, - DeviceModule, - EnergyModule, - LedModule, - LightTransitionModule, - TimeModule, -) -from .smartmodule import SmartModule +from .modules import * # noqa: F403 + _LOGGER = logging.getLogger(__name__) if TYPE_CHECKING: - from .smartchilddevice import SmartChildDevice + from .smartmodule import SmartModule class SmartDevice(Device): @@ -49,21 +40,32 @@ def __init__( self._components: Dict[str, int] = {} self._children: Dict[str, "SmartChildDevice"] = {} self._state_information: Dict[str, Any] = {} - self.modules: Dict[str, SmartModule] = {} + self.modules: Dict[str, "SmartModule"] = {} + self._parent: Optional["SmartDevice"] = None async def _initialize_children(self): """Initialize children for power strips.""" - children = self._last_update["child_info"]["child_device_list"] - # TODO: Use the type information to construct children, - # as hubs can also have them. + children = self.internal_state["child_info"]["child_device_list"] + children_components = { + child["device_id"]: { + comp["id"]: int(comp["ver_code"]) for comp in child["component_list"] + } + for child in self.internal_state["get_child_device_component_list"][ + "child_component_list" + ] + } from .smartchilddevice import SmartChildDevice self._children = { - child["device_id"]: SmartChildDevice( - parent=self, child_id=child["device_id"] + child_info["device_id"]: await SmartChildDevice.create( + parent=self, + child_info=child_info, + child_components=children_components[child_info["device_id"]], ) - for child in children + for child_info in children } + # TODO: if all are sockets, then we are a strip, and otherwise a hub? + # doesn't work for the walldimmer with fancontrol... self._device_type = DeviceType.Strip @property @@ -126,8 +128,10 @@ async def update(self, update_children: bool = True): if not self.children: await self._initialize_children() + # TODO: we don't currently perform queries on children based on modules, + # but just update the information that is returned in the main query. for info in child_info["child_device_list"]: - self._children[info["device_id"]].update_internal_state(info) + 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. @@ -153,6 +157,7 @@ async def _initialize_modules(self): async def _initialize_features(self): """Initialize device features.""" + self._add_feature(Feature(self, "Device ID", attribute_getter="device_id")) if "device_on" in self._info: self._add_feature( Feature( @@ -164,25 +169,33 @@ async def _initialize_features(self): ) ) - self._add_feature( - Feature( - self, - "Signal Level", - attribute_getter=lambda x: x._info["signal_level"], - icon="mdi:signal", + if "signal_level" in self._info: + + 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", + + if "rssi" in self._info: + self._add_feature( + Feature( + self, + "RSSI", + attribute_getter=lambda x: x._info["rssi"], + icon="mdi:signal", + ) + ) + + if "ssid" in self._info: + self._add_feature( + Feature( + device=self, name="SSID", attribute_getter="ssid", icon="mdi:wifi" + ) ) - ) - self._add_feature( - Feature(device=self, name="SSID", attribute_getter="ssid", icon="mdi:wifi") - ) if "overheated" in self._info: self._add_feature( @@ -232,7 +245,12 @@ def alias(self) -> Optional[str]: @property def time(self) -> datetime: """Return the time.""" - _timemod = cast(TimeModule, self.modules["TimeModule"]) + # TODO: Default to parent's time module for child devices + if self._parent and "TimeModule" in self.modules: + _timemod = cast(TimeModule, self._parent.modules["TimeModule"]) # noqa: F405 + else: + _timemod = cast(TimeModule, self.modules["TimeModule"]) # noqa: F405 + return _timemod.time @property @@ -284,6 +302,14 @@ def internal_state(self) -> Any: """Return all the internal state data.""" return self._last_update + def _update_internal_state(self, info): + """Update internal state. + + This is used by the parent to push updates to its children + """ + # TODO: cleanup the _last_update, _info mess. + self._last_update = self._info = info + async def _query_helper( self, method: str, params: Optional[Dict] = None, child_ids=None ) -> Any: @@ -347,19 +373,19 @@ async def get_emeter_realtime(self) -> EmeterStatus: @property def emeter_realtime(self) -> EmeterStatus: """Get the emeter status.""" - energy = cast(EnergyModule, self.modules["EnergyModule"]) + energy = cast(EnergyModule, self.modules["EnergyModule"]) # noqa: F405 return energy.emeter_realtime @property def emeter_this_month(self) -> Optional[float]: """Get the emeter value for this month.""" - energy = cast(EnergyModule, self.modules["EnergyModule"]) + energy = cast(EnergyModule, self.modules["EnergyModule"]) # noqa: F405 return energy.emeter_this_month @property def emeter_today(self) -> Optional[float]: """Get the emeter value for today.""" - energy = cast(EnergyModule, self.modules["EnergyModule"]) + energy = cast(EnergyModule, self.modules["EnergyModule"]) # noqa: F405 return energy.emeter_today @property @@ -372,7 +398,7 @@ def on_since(self) -> Optional[datetime]: return None on_time = cast(float, on_time) if (timemod := self.modules.get("TimeModule")) is not None: - timemod = cast(TimeModule, timemod) + timemod = cast(TimeModule, timemod) # noqa: F405 return timemod.time - timedelta(seconds=on_time) else: # We have no device time, use current local time. return datetime.now().replace(microsecond=0) - timedelta(seconds=on_time) diff --git a/kasa/smart/smartmodule.py b/kasa/smart/smartmodule.py index 6f42f297a..14e0d92bd 100644 --- a/kasa/smart/smartmodule.py +++ b/kasa/smart/smartmodule.py @@ -57,15 +57,25 @@ def data(self): """ q = self.query() q_keys = list(q.keys()) + query_key = q_keys[0] + + dev = self._device + # TODO: hacky way to check if update has been called. - if q_keys[0] not in self._device._last_update: - raise SmartDeviceException( - f"You need to call update() prior accessing module data" - f" for '{self._module}'" - ) + # The way this falls back to parent may not always be wanted. + # Especially, devices can have their own firmware updates. + if query_key not in dev._last_update: + if dev._parent and query_key in dev._parent._last_update: + _LOGGER.debug("%s not found child, but found on parent", query_key) + dev = dev._parent + else: + raise SmartDeviceException( + f"You need to call update() prior accessing module data" + f" for '{self._module}'" + ) filtered_data = { - k: v for k, v in self._device._last_update.items() if k in q_keys + k: v for k, v in dev._last_update.items() if k in q_keys } if len(filtered_data) == 1: return next(iter(filtered_data.values())) diff --git a/kasa/tests/test_childdevice.py b/kasa/tests/test_childdevice.py index 78863def3..6ffd70549 100644 --- a/kasa/tests/test_childdevice.py +++ b/kasa/tests/test_childdevice.py @@ -47,7 +47,6 @@ async def test_childdevice_properties(dev: SmartChildDevice): assert len(dev.children) > 0 first = dev.children[0] - assert first.is_strip_socket # children do not have children assert not first.children diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py index 487286dbb..61659945c 100644 --- a/kasa/tests/test_smartdevice.py +++ b/kasa/tests/test_smartdevice.py @@ -255,7 +255,9 @@ async def test_device_class_ctors(device_class_name_obj): klass = device_class_name_obj[1] if issubclass(klass, SmartChildDevice): parent = SmartDevice(host, config=config) - dev = klass(parent, 1) + dev = klass( + parent, {"dummy": "info", "device_id": "dummy"}, {"dummy": "components"} + ) else: dev = klass(host, config=config) assert dev.host == host From df6b95c39a89bd9df6e56da33787b0a3d67ceb2b Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Tue, 20 Feb 2024 01:38:45 +0100 Subject: [PATCH 2/8] Fix linting --- kasa/cli.py | 5 +++-- kasa/smart/modules/__init__.py | 1 - kasa/smart/smartdevice.py | 3 +-- kasa/smart/smartmodule.py | 4 +--- 4 files changed, 5 insertions(+), 8 deletions(-) diff --git a/kasa/cli.py b/kasa/cli.py index 0a2fff8a1..2e3117ae5 100755 --- a/kasa/cli.py +++ b/kasa/cli.py @@ -656,9 +656,10 @@ async def cmd_command(dev: Device, module, child, command, parameters): if child: echo(f"Selecting child {child} from {dev}") - # TODO: Update required to initialize children + # TODO: Update required to initialize children, await dev.update() - dev = dev._children[child] + # TODO: _children is not yet part of the Device API. + dev = dev._children[child] # type: ignore[attr-defined] if isinstance(dev, IotDevice): res = await dev._query_helper(module, command, parameters) diff --git a/kasa/smart/modules/__init__.py b/kasa/smart/modules/__init__.py index c8e17c57f..02c3b86af 100644 --- a/kasa/smart/modules/__init__.py +++ b/kasa/smart/modules/__init__.py @@ -8,7 +8,6 @@ from .ledmodule import LedModule from .lighttransitionmodule import LightTransitionModule from .timemodule import TimeModule -from .firmware import Firmware __all__ = [ "TimeModule", diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py index 9e691c774..b41386f59 100644 --- a/kasa/smart/smartdevice.py +++ b/kasa/smart/smartdevice.py @@ -14,10 +14,10 @@ from ..smartprotocol import SmartProtocol from .modules import * # noqa: F403 - _LOGGER = logging.getLogger(__name__) if TYPE_CHECKING: + from .smartchilddevice import SmartChildDevice from .smartmodule import SmartModule @@ -170,7 +170,6 @@ async def _initialize_features(self): ) if "signal_level" in self._info: - self._add_feature( Feature( self, diff --git a/kasa/smart/smartmodule.py b/kasa/smart/smartmodule.py index 14e0d92bd..097e6ff81 100644 --- a/kasa/smart/smartmodule.py +++ b/kasa/smart/smartmodule.py @@ -74,9 +74,7 @@ def data(self): f" for '{self._module}'" ) - filtered_data = { - k: v for k, v in dev._last_update.items() if k in q_keys - } + filtered_data = {k: v for k, v in dev._last_update.items() if k in q_keys} if len(filtered_data) == 1: return next(iter(filtered_data.values())) From d28984ecd7e4372df6e1dd0cc02cdd3967aa9a45 Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Wed, 21 Feb 2024 15:59:18 +0100 Subject: [PATCH 3/8] Add comment why we are calling update() for raw commands on children --- kasa/cli.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/kasa/cli.py b/kasa/cli.py index 2e3117ae5..481c126b9 100755 --- a/kasa/cli.py +++ b/kasa/cli.py @@ -655,10 +655,11 @@ async def cmd_command(dev: Device, module, child, command, parameters): parameters = ast.literal_eval(parameters) if child: - echo(f"Selecting child {child} from {dev}") - # TODO: Update required to initialize children, + # The way child devices are accessed requires a ChildDevice to + # wrap the communications. Doing this properly would require creating + # a common interfaces for both IOT and SMART child devices. + # As a stop-gap solution, we perform an update instead. await dev.update() - # TODO: _children is not yet part of the Device API. dev = dev._children[child] # type: ignore[attr-defined] if isinstance(dev, IotDevice): From 73be0d5854b080df1c24f09845fce881d7d9f750 Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Wed, 21 Feb 2024 16:42:45 +0100 Subject: [PATCH 4/8] Add get_child_device to base class Also, pull _children up to the base class and make necessary changes to make it work. Remove has_children from iotdevice, which was only used in tests. --- kasa/cli.py | 2 +- kasa/device.py | 28 +++++++++++++++++++--------- kasa/iot/iotdevice.py | 17 +---------------- kasa/iot/iotstrip.py | 6 +++--- kasa/smart/smartdevice.py | 2 -- kasa/tests/test_smartdevice.py | 2 -- 6 files changed, 24 insertions(+), 33 deletions(-) diff --git a/kasa/cli.py b/kasa/cli.py index 481c126b9..747008bc9 100755 --- a/kasa/cli.py +++ b/kasa/cli.py @@ -660,7 +660,7 @@ async def cmd_command(dev: Device, module, child, command, parameters): # a common interfaces for both IOT and SMART child devices. # As a stop-gap solution, we perform an update instead. await dev.update() - dev = dev._children[child] # type: ignore[attr-defined] + dev = dev.get_child_device(child) if isinstance(dev, IotDevice): res = await dev._query_helper(module, command, parameters) diff --git a/kasa/device.py b/kasa/device.py index fb92f18e9..39dd6bff1 100644 --- a/kasa/device.py +++ b/kasa/device.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from datetime import datetime -from typing import Any, Dict, List, Optional, Sequence, Union +from typing import Any, Dict, Generic, List, Optional, Sequence, TypeVar, Union from .credentials import Credentials from .device_type import DeviceType @@ -35,7 +35,10 @@ class WifiNetwork: _LOGGER = logging.getLogger(__name__) -class Device(ABC): +T = TypeVar("T", bound="Device") + + +class Device(ABC, Generic[T]): """Common device interface. Do not instantiate this class directly, instead get a device instance from @@ -71,14 +74,15 @@ def __init__( self.modules: Dict[str, Any] = {} self._features: Dict[str, Feature] = {} - self._parent: Optional["Device"] = None + self._parent: Optional[T] = None + self._children: Dict[str, T] = {} @staticmethod async def connect( *, host: Optional[str] = None, config: Optional[DeviceConfig] = None, - ) -> "Device": + ) -> T: """Connect to a single device by the given hostname or device configuration. This method avoids the UDP based discovery process and @@ -97,7 +101,9 @@ async def connect( """ from .device_factory import connect # pylint: disable=import-outside-toplevel - return await connect(host=host, config=config) # type: ignore[arg-type] + # TODO: Incompatible return value type (got "Device[Any]", expected "T") + # unclear how to type that properly for now. + return await connect(host=host, config=config) # type: ignore[arg-type,return-value] @abstractmethod async def update(self, update_children: bool = True): @@ -183,9 +189,13 @@ async def _raw_query(self, request: Union[str, Dict]) -> Any: return await self.protocol.query(request=request) @property - @abstractmethod - def children(self) -> Sequence["Device"]: + def children(self) -> Sequence[T]: """Returns the child devices.""" + return list(self._children.values()) + + def get_child_device(self, id_: str): + """Return child device by its ID.""" + return self._children[id_] @property @abstractmethod @@ -237,7 +247,7 @@ def is_color(self) -> bool: """Return True if the device supports color changes.""" return False - def get_plug_by_name(self, name: str) -> "Device": + def get_plug_by_name(self, name: str) -> T: """Return child device for the given name.""" for p in self.children: if p.alias == name: @@ -245,7 +255,7 @@ def get_plug_by_name(self, name: str) -> "Device": raise SmartDeviceException(f"Device has no child with {name}") - def get_plug_by_index(self, index: int) -> "Device": + def get_plug_by_index(self, index: int) -> T: """Return child device for the given index.""" if index + 1 > len(self.children) or index < 0: raise SmartDeviceException( diff --git a/kasa/iot/iotdevice.py b/kasa/iot/iotdevice.py index ac902af84..36fcb2547 100755 --- a/kasa/iot/iotdevice.py +++ b/kasa/iot/iotdevice.py @@ -186,19 +186,13 @@ def __init__( super().__init__(host=host, config=config, protocol=protocol) self._sys_info: Any = None # TODO: this is here to avoid changing tests - self._children: Sequence["IotDevice"] = [] self._supported_modules: Optional[Dict[str, IotModule]] = None self._legacy_features: Set[str] = set() @property def children(self) -> Sequence["IotDevice"]: """Return list of children.""" - return self._children - - @children.setter - def children(self, children): - """Initialize from a list of children.""" - self._children = children + return list(self._children.values()) def add_module(self, name: str, module: IotModule): """Register a module.""" @@ -411,15 +405,6 @@ def model(self) -> str: sys_info = self._sys_info return str(sys_info["model"]) - @property - def has_children(self) -> bool: - """Return true if the device has children devices.""" - # Ideally we would check for the 'child_num' key in sys_info, - # but devices that speak klap do not populate this key via - # update_from_discover_info so we check for the devices - # we know have children instead. - return self.is_strip - @property # type: ignore def alias(self) -> Optional[str]: """Return device name (alias).""" diff --git a/kasa/iot/iotstrip.py b/kasa/iot/iotstrip.py index 7cbb10b03..2f2e2beb2 100755 --- a/kasa/iot/iotstrip.py +++ b/kasa/iot/iotstrip.py @@ -115,10 +115,10 @@ async def update(self, update_children: bool = True): if not self.children: children = self.sys_info["children"] _LOGGER.debug("Initializing %s child sockets", len(children)) - self.children = [ - IotStripPlug(self.host, parent=self, child_id=child["id"]) + self._children = { + child["id"]: IotStripPlug(self.host, parent=self, child_id=child["id"]) for child in children - ] + } if update_children and self.has_emeter: for plug in self.children: diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py index b41386f59..63f2ebe2d 100644 --- a/kasa/smart/smartdevice.py +++ b/kasa/smart/smartdevice.py @@ -17,7 +17,6 @@ _LOGGER = logging.getLogger(__name__) if TYPE_CHECKING: - from .smartchilddevice import SmartChildDevice from .smartmodule import SmartModule @@ -38,7 +37,6 @@ def __init__( self.protocol: SmartProtocol self._components_raw: Optional[Dict[str, Any]] = None self._components: Dict[str, int] = {} - self._children: Dict[str, "SmartChildDevice"] = {} self._state_information: Dict[str, Any] = {} self.modules: Dict[str, "SmartModule"] = {} self._parent: Optional["SmartDevice"] = None diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py index 61659945c..4fa42bc01 100644 --- a/kasa/tests/test_smartdevice.py +++ b/kasa/tests/test_smartdevice.py @@ -210,10 +210,8 @@ async def test_children(dev): """Make sure that children property is exposed by every device.""" if dev.is_strip: assert len(dev.children) > 0 - assert dev.has_children is True else: assert len(dev.children) == 0 - assert dev.has_children is False @device_iot From 446cbcebcddc78e5d5483d7980df20c7e522889b Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Thu, 22 Feb 2024 19:09:08 +0100 Subject: [PATCH 5/8] Remove generics --- kasa/device.py | 21 +++++++++------------ kasa/iot/iotdevice.py | 3 ++- kasa/smart/smartdevice.py | 3 ++- 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/kasa/device.py b/kasa/device.py index 39dd6bff1..6ce419aa7 100644 --- a/kasa/device.py +++ b/kasa/device.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from datetime import datetime -from typing import Any, Dict, Generic, List, Optional, Sequence, TypeVar, Union +from typing import Any, Dict, List, Mapping, Optional, Sequence, Union from .credentials import Credentials from .device_type import DeviceType @@ -35,10 +35,7 @@ class WifiNetwork: _LOGGER = logging.getLogger(__name__) -T = TypeVar("T", bound="Device") - - -class Device(ABC, Generic[T]): +class Device(ABC): """Common device interface. Do not instantiate this class directly, instead get a device instance from @@ -74,15 +71,15 @@ def __init__( self.modules: Dict[str, Any] = {} self._features: Dict[str, Feature] = {} - self._parent: Optional[T] = None - self._children: Dict[str, T] = {} + self._parent: Optional["Device"] = None + self._children: Mapping[str, "Device"] = {} @staticmethod async def connect( *, host: Optional[str] = None, config: Optional[DeviceConfig] = None, - ) -> T: + ) -> "Device": """Connect to a single device by the given hostname or device configuration. This method avoids the UDP based discovery process and @@ -189,11 +186,11 @@ async def _raw_query(self, request: Union[str, Dict]) -> Any: return await self.protocol.query(request=request) @property - def children(self) -> Sequence[T]: + def children(self) -> Sequence["Device"]: """Returns the child devices.""" return list(self._children.values()) - def get_child_device(self, id_: str): + def get_child_device(self, id_: str) -> "Device": """Return child device by its ID.""" return self._children[id_] @@ -247,7 +244,7 @@ def is_color(self) -> bool: """Return True if the device supports color changes.""" return False - def get_plug_by_name(self, name: str) -> T: + def get_plug_by_name(self, name: str) -> "Device": """Return child device for the given name.""" for p in self.children: if p.alias == name: @@ -255,7 +252,7 @@ def get_plug_by_name(self, name: str) -> T: raise SmartDeviceException(f"Device has no child with {name}") - def get_plug_by_index(self, index: int) -> T: + def get_plug_by_index(self, index: int) -> "Device": """Return child device for the given index.""" if index + 1 > len(self.children) or index < 0: raise SmartDeviceException( diff --git a/kasa/iot/iotdevice.py b/kasa/iot/iotdevice.py index 36fcb2547..3bef4def1 100755 --- a/kasa/iot/iotdevice.py +++ b/kasa/iot/iotdevice.py @@ -16,7 +16,7 @@ import inspect import logging from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional, Sequence, Set +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set from ..device import Device, WifiNetwork from ..deviceconfig import DeviceConfig @@ -188,6 +188,7 @@ def __init__( self._sys_info: Any = None # TODO: this is here to avoid changing tests self._supported_modules: Optional[Dict[str, IotModule]] = None self._legacy_features: Set[str] = set() + self._children: Mapping[str, "IotDevice"] = {} @property def children(self) -> Sequence["IotDevice"]: diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py index 63f2ebe2d..b8bd65384 100644 --- a/kasa/smart/smartdevice.py +++ b/kasa/smart/smartdevice.py @@ -2,7 +2,7 @@ import base64 import logging from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, cast +from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Sequence, cast from ..aestransport import AesTransport from ..device import Device, WifiNetwork @@ -40,6 +40,7 @@ def __init__( self._state_information: Dict[str, Any] = {} self.modules: Dict[str, "SmartModule"] = {} self._parent: Optional["SmartDevice"] = None + self._children: Mapping[str, "SmartDevice"] = {} async def _initialize_children(self): """Initialize children for power strips.""" From 26640f09538467369cf3f5c6a1fe5c75c977a1b7 Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Thu, 22 Feb 2024 19:13:36 +0100 Subject: [PATCH 6/8] Remove unnecessary type ignore --- kasa/device.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/kasa/device.py b/kasa/device.py index 6ce419aa7..5080ddb07 100644 --- a/kasa/device.py +++ b/kasa/device.py @@ -98,9 +98,7 @@ async def connect( """ from .device_factory import connect # pylint: disable=import-outside-toplevel - # TODO: Incompatible return value type (got "Device[Any]", expected "T") - # unclear how to type that properly for now. - return await connect(host=host, config=config) # type: ignore[arg-type,return-value] + return await connect(host=host, config=config) # type: ignore[arg-type] @abstractmethod async def update(self, update_children: bool = True): From 6f8f730c36a0f951de9d37b8a234d3a766a1faee Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Thu, 22 Feb 2024 20:37:26 +0100 Subject: [PATCH 7/8] Add tests, fix iotstrip child_id issue --- kasa/iot/iotstrip.py | 2 +- kasa/tests/test_cli.py | 24 ++++++++++++++++++++++++ kasa/tests/test_smartdevice.py | 15 +++++++-------- 3 files changed, 32 insertions(+), 9 deletions(-) diff --git a/kasa/iot/iotstrip.py b/kasa/iot/iotstrip.py index 15dbf906e..59ea71693 100755 --- a/kasa/iot/iotstrip.py +++ b/kasa/iot/iotstrip.py @@ -116,7 +116,7 @@ async def update(self, update_children: bool = True): children = self.sys_info["children"] _LOGGER.debug("Initializing %s child sockets", len(children)) self._children = { - child["id"]: IotStripPlug(self.host, parent=self, child_id=child["id"]) + f"{self.mac}_{child['id']}": IotStripPlug(self.host, parent=self, child_id=child["id"]) for child in children } diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py index 8bffef7d6..9c9ebed9e 100644 --- a/kasa/tests/test_cli.py +++ b/kasa/tests/test_cli.py @@ -21,6 +21,7 @@ cli, emeter, raw_command, + cmd_command, reboot, state, sysinfo, @@ -136,6 +137,29 @@ async def test_raw_command(dev, mocker): assert "Usage" in res.output +async def test_command_with_child(dev, mocker): + """Test 'command' command with --child.""" + runner = CliRunner() + update_mock = mocker.patch.object(dev, "update") + + dummy_child = mocker.create_autospec(IotDevice) + query_mock = mocker.patch.object(dummy_child, "_query_helper", return_value={"dummy": "response"}) + + mocker.patch.object(dev, "_children", {"XYZ": dummy_child}) + get_child_device_mock = mocker.patch.object(dev, "get_child_device", return_value=dummy_child) + + res = await runner.invoke( + cmd_command, + ["--child", "XYZ", "command", "'params'"], + obj=dev, catch_exceptions=False + ) + + update_mock.assert_called() + query_mock.assert_called() + assert '{"dummy": "response"}' in res.output + assert res.exit_code == 0 + + @device_smart async def test_reboot(dev, mocker): """Test that reboot works on SMART devices.""" diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py index d0a09b9ba..ec83255fe 100644 --- a/kasa/tests/test_smartdevice.py +++ b/kasa/tests/test_smartdevice.py @@ -38,6 +38,7 @@ no_emeter_iot, plug, turn_on, + strip, ) from .fakeprotocol_iot import FakeIotProtocol @@ -201,14 +202,12 @@ async def test_representation(dev): assert pattern.match(str(dev)) -@device_iot -async def test_childrens(dev): - """Make sure that children property is exposed by every device.""" - if dev.is_strip: - assert len(dev.children) > 0 - else: - assert len(dev.children) == 0 - +@strip +def test_children_api(dev): + """Make sure the children api""" + first = dev.children[0] + first_by_get_child_device = dev.get_child_device(first.device_id) + assert first == first_by_get_child_device @device_iot async def test_children(dev): From 6eb8de4b464d7dac49016b11a48fb9ec86ce2fed Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Thu, 22 Feb 2024 20:41:55 +0100 Subject: [PATCH 8/8] Fix linting --- kasa/iot/iotstrip.py | 4 +++- kasa/tests/test_cli.py | 11 +++++++---- kasa/tests/test_smartdevice.py | 5 +++-- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/kasa/iot/iotstrip.py b/kasa/iot/iotstrip.py index 59ea71693..4bf31cc76 100755 --- a/kasa/iot/iotstrip.py +++ b/kasa/iot/iotstrip.py @@ -116,7 +116,9 @@ async def update(self, update_children: bool = True): children = self.sys_info["children"] _LOGGER.debug("Initializing %s child sockets", len(children)) self._children = { - f"{self.mac}_{child['id']}": IotStripPlug(self.host, parent=self, child_id=child["id"]) + f"{self.mac}_{child['id']}": IotStripPlug( + self.host, parent=self, child_id=child["id"] + ) for child in children } diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py index 9c9ebed9e..4c0d17e13 100644 --- a/kasa/tests/test_cli.py +++ b/kasa/tests/test_cli.py @@ -19,9 +19,9 @@ alias, brightness, cli, + cmd_command, emeter, raw_command, - cmd_command, reboot, state, sysinfo, @@ -143,15 +143,18 @@ async def test_command_with_child(dev, mocker): update_mock = mocker.patch.object(dev, "update") dummy_child = mocker.create_autospec(IotDevice) - query_mock = mocker.patch.object(dummy_child, "_query_helper", return_value={"dummy": "response"}) + query_mock = mocker.patch.object( + dummy_child, "_query_helper", return_value={"dummy": "response"} + ) mocker.patch.object(dev, "_children", {"XYZ": dummy_child}) - get_child_device_mock = mocker.patch.object(dev, "get_child_device", return_value=dummy_child) + mocker.patch.object(dev, "get_child_device", return_value=dummy_child) res = await runner.invoke( cmd_command, ["--child", "XYZ", "command", "'params'"], - obj=dev, catch_exceptions=False + obj=dev, + catch_exceptions=False, ) update_mock.assert_called() diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py index ec83255fe..92cca5a16 100644 --- a/kasa/tests/test_smartdevice.py +++ b/kasa/tests/test_smartdevice.py @@ -37,8 +37,8 @@ lightstrip, no_emeter_iot, plug, - turn_on, strip, + turn_on, ) from .fakeprotocol_iot import FakeIotProtocol @@ -204,11 +204,12 @@ async def test_representation(dev): @strip def test_children_api(dev): - """Make sure the children api""" + """Test the child device API.""" first = dev.children[0] first_by_get_child_device = dev.get_child_device(first.device_id) assert first == first_by_get_child_device + @device_iot async def test_children(dev): """Make sure that children property is exposed by every device."""