From 418e5088add78e96f60550ee8e5dddf81f38bdf2 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 19 Nov 2023 09:56:23 -0600 Subject: [PATCH] move method --- kasa/device_factory.py | 101 +++++++++++++++++++++++++++++++++++++++++ kasa/device_type.py | 25 ++++++++++ kasa/discover.py | 92 ++----------------------------------- kasa/smartdevice.py | 62 ++++++++++++++++--------- 4 files changed, 170 insertions(+), 110 deletions(-) create mode 100755 kasa/device_factory.py create mode 100755 kasa/device_type.py diff --git a/kasa/device_factory.py b/kasa/device_factory.py new file mode 100755 index 00000000..a18b0912 --- /dev/null +++ b/kasa/device_factory.py @@ -0,0 +1,101 @@ +"""Device creation by type.""" + +from typing import Optional, Type + +from .credentials import Credentials +from .device_type import DeviceType +from .smartbulb import SmartBulb +from .smartdevice import SmartDevice, SmartDeviceException +from .smartdimmer import SmartDimmer +from .smartlightstrip import SmartLightStrip +from .smartplug import SmartPlug +from .smartstrip import SmartStrip + +DEVICE_TYPE_TO_CLASS = { + DeviceType.Plug: SmartPlug, + DeviceType.Bulb: SmartBulb, + DeviceType.Strip: SmartStrip, + DeviceType.Dimmer: SmartDimmer, + DeviceType.LightStrip: SmartLightStrip, +} + + +def get_device_class_from_type(device_type: DeviceType) -> Type[SmartDevice]: + """Find SmartDevice subclass for device described by passed data.""" + return DEVICE_TYPE_TO_CLASS[device_type] + + +async def connect( + host: str, + *, + port: Optional[int] = None, + timeout=5, + credentials: Optional[Credentials] = None, + device_type: Optional[DeviceType] = None, +) -> "SmartDevice": + """Connect to a single device by the given IP address. + + This method avoids the UDP based discovery process and + will connect directly to the device to query its type. + + It is generally preferred to avoid :func:`discover_single()` and + use this function instead as it should perform better when + the WiFi network is congested or the device is not responding + to discovery requests. + + The device type is discovered by querying the device. + + :param host: Hostname of device to query + :param device_type: Device type to use for the device. + If not given, the device type is discovered by querying the device. + If the device type is already known, it is preferred to pass it + to avoid the extra query to the device to discover its type. + :rtype: SmartDevice + :return: Object for querying/controlling found device. + """ + if device_type and (klass := DEVICE_TYPE_TO_CLASS.get(device_type)): + dev: SmartDevice = klass( + host=host, port=port, credentials=credentials, timeout=timeout + ) + await dev.update() + return dev + + unknown_dev = SmartDevice( + host=host, port=port, credentials=credentials, timeout=timeout + ) + await unknown_dev.update() + device_class = get_device_class_from_info(unknown_dev.internal_state) + dev = device_class(host=host, port=port, credentials=credentials, timeout=timeout) + # Reuse the connection from the unknown device + # so we don't have to reconnect + dev.protocol = unknown_dev.protocol + await dev.update() + return dev + + +def get_device_class_from_info(info: dict) -> Type[SmartDevice]: + """Find SmartDevice subclass for device described by passed data.""" + if "system" not in info or "get_sysinfo" not in info["system"]: + raise SmartDeviceException("No 'system' or 'get_sysinfo' in response") + + sysinfo = info["system"]["get_sysinfo"] + type_ = sysinfo.get("type", sysinfo.get("mic_type")) + if type_ is None: + raise SmartDeviceException("Unable to find the device type field!") + + if "dev_name" in sysinfo and "Dimmer" in sysinfo["dev_name"]: + return SmartDimmer + + if "smartplug" in type_.lower(): + if "children" in sysinfo: + return SmartStrip + + return SmartPlug + + if "smartbulb" in type_.lower(): + if "length" in sysinfo: # strips have length + return SmartLightStrip + + return SmartBulb + + raise SmartDeviceException("Unknown device type: %s" % type_) diff --git a/kasa/device_type.py b/kasa/device_type.py new file mode 100755 index 00000000..162fc4f2 --- /dev/null +++ b/kasa/device_type.py @@ -0,0 +1,25 @@ +"""TP-Link device types.""" + + +from enum import Enum + + +class DeviceType(Enum): + """Device type enum.""" + + # The values match what the cli has historically used + Plug = "plug" + Bulb = "bulb" + Strip = "strip" + StripSocket = "stripsocket" + Dimmer = "dimmer" + LightStrip = "lightstrip" + Unknown = "unknown" + + @staticmethod + def from_value(name: str) -> "DeviceType": + """Return device type from string value.""" + for device_type in DeviceType: + if device_type.value == name: + return device_type + return DeviceType.Unknown diff --git a/kasa/discover.py b/kasa/discover.py index db7235b1..2580b699 100755 --- a/kasa/discover.py +++ b/kasa/discover.py @@ -15,12 +15,9 @@ from kasa.exceptions import UnsupportedDeviceException from kasa.json import dumps as json_dumps from kasa.json import loads as json_loads from kasa.protocol import TPLinkSmartHomeProtocol -from kasa.smartbulb import SmartBulb -from kasa.smartdevice import DeviceType, SmartDevice, SmartDeviceException -from kasa.smartdimmer import SmartDimmer -from kasa.smartlightstrip import SmartLightStrip -from kasa.smartplug import SmartPlug -from kasa.smartstrip import SmartStrip +from kasa.smartdevice import SmartDevice, SmartDeviceException + +from .device_factory import get_device_class_from_info _LOGGER = logging.getLogger(__name__) @@ -28,14 +25,6 @@ _LOGGER = logging.getLogger(__name__) OnDiscoveredCallable = Callable[[SmartDevice], Awaitable[None]] DeviceDict = Dict[str, SmartDevice] -DEVICE_TYPE_TO_CLASS = { - DeviceType.Plug: SmartPlug, - DeviceType.Bulb: SmartBulb, - DeviceType.Strip: SmartStrip, - DeviceType.Dimmer: SmartDimmer, - DeviceType.LightStrip: SmartLightStrip, -} - class _DiscoverProtocol(asyncio.DatagramProtocol): """Implementation of the discovery protocol handler. @@ -345,80 +334,7 @@ class Discover: else: raise SmartDeviceException(f"Unable to get discovery response for {host}") - @staticmethod - async def connect_single( - host: str, - *, - port: Optional[int] = None, - timeout=5, - credentials: Optional[Credentials] = None, - device_type: Optional[DeviceType] = None, - ) -> SmartDevice: - """Connect to a single device by the given IP address. - - This method avoids the UDP based discovery process and - will connect directly to the device to query its type. - - It is generally preferred to avoid :func:`discover_single()` and - use this function instead as it should perform better when - the WiFi network is congested or the device is not responding - to discovery requests. - - The device type is discovered by querying the device. - - :param host: Hostname of device to query - :param device_type: Device type to use for the device. - If not given, the device type is discovered by querying the device. - If the device type is already known, it is preferred to pass it - to avoid the extra query to the device to discover its type. - :rtype: SmartDevice - :return: Object for querying/controlling found device. - """ - if device_type and (klass := DEVICE_TYPE_TO_CLASS.get(device_type)): - dev: SmartDevice = klass( - host=host, port=port, credentials=credentials, timeout=timeout - ) - await dev.update() - return dev - - unknown_dev = SmartDevice( - host=host, port=port, credentials=credentials, timeout=timeout - ) - await unknown_dev.update() - device_class = Discover._get_device_class(unknown_dev.internal_state) - dev = device_class( - host=host, port=port, credentials=credentials, timeout=timeout - ) - # Reuse the connection from the unknown device - # so we don't have to reconnect - dev.protocol = unknown_dev.protocol - await dev.update() - return dev - @staticmethod def _get_device_class(info: dict) -> Type[SmartDevice]: """Find SmartDevice subclass for device described by passed data.""" - if "system" not in info or "get_sysinfo" not in info["system"]: - raise SmartDeviceException("No 'system' or 'get_sysinfo' in response") - - sysinfo = info["system"]["get_sysinfo"] - type_ = sysinfo.get("type", sysinfo.get("mic_type")) - if type_ is None: - raise SmartDeviceException("Unable to find the device type field!") - - if "dev_name" in sysinfo and "Dimmer" in sysinfo["dev_name"]: - return SmartDimmer - - if "smartplug" in type_.lower(): - if "children" in sysinfo: - return SmartStrip - - return SmartPlug - - if "smartbulb" in type_.lower(): - if "length" in sysinfo: # strips have length - return SmartLightStrip - - return SmartBulb - - raise SmartDeviceException("Unknown device type: %s" % type_) + return get_device_class_from_info(info) diff --git a/kasa/smartdevice.py b/kasa/smartdevice.py index e8142a5d..4f850b5b 100755 --- a/kasa/smartdevice.py +++ b/kasa/smartdevice.py @@ -17,10 +17,10 @@ import inspect import logging from dataclasses import dataclass from datetime import datetime, timedelta -from enum import Enum from typing import Any, Dict, List, Optional, Set from .credentials import Credentials +from .device_type import DeviceType from .emeterstatus import EmeterStatus from .exceptions import SmartDeviceException from .modules import Emeter, Module @@ -29,27 +29,6 @@ from .protocol import TPLinkSmartHomeProtocol _LOGGER = logging.getLogger(__name__) -class DeviceType(Enum): - """Device type enum.""" - - # The values match what the cli has historically used - Plug = "plug" - Bulb = "bulb" - Strip = "strip" - StripSocket = "stripsocket" - Dimmer = "dimmer" - LightStrip = "lightstrip" - Unknown = "unknown" - - @staticmethod - def from_value(name: str) -> "DeviceType": - """Return device type from string value.""" - for device_type in DeviceType: - if device_type.value == name: - return device_type - return DeviceType.Unknown - - @dataclass class WifiNetwork: """Wifi network container.""" @@ -767,3 +746,42 @@ class SmartDevice: f" ({self.alias}), is_on: {self.is_on}" f" - dev specific: {self.state_information}>" ) + + @staticmethod + async def connect( + host: str, + *, + port: Optional[int] = None, + timeout=5, + credentials: Optional[Credentials] = None, + device_type: Optional[DeviceType] = None, + ) -> "SmartDevice": + """Connect to a single device by the given IP address. + + This method avoids the UDP based discovery process and + will connect directly to the device to query its type. + + It is generally preferred to avoid :func:`discover_single()` and + use this function instead as it should perform better when + the WiFi network is congested or the device is not responding + to discovery requests. + + The device type is discovered by querying the device. + + :param host: Hostname of device to query + :param device_type: Device type to use for the device. + If not given, the device type is discovered by querying the device. + If the device type is already known, it is preferred to pass it + to avoid the extra query to the device to discover its type. + :rtype: SmartDevice + :return: Object for querying/controlling found device. + """ + from .device_factory import connect # pylint: disable=import-outside-toplevel + + return await connect( + host=host, + port=port, + timeout=timeout, + credentials=credentials, + device_type=device_type, + )