move method

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_)

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

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 SmartDevice, SmartDeviceException
-from kasa.smartdevice import DeviceType, SmartDevice, SmartDeviceException
+
-from kasa.smartdimmer import SmartDimmer
+from .device_factory import get_device_class_from_info
-from kasa.smartlightstrip import SmartLightStrip
-from kasa.smartplug import SmartPlug
-from kasa.smartstrip import SmartStrip
 
 _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"]:
+        return get_device_class_from_info(info)
-            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_)

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,
+        )