python-kasa · rytilahti · Nov 21, 2023 · Oct 31, 2023 · Oct 31, 2023 · Nov 1, 2023
diff --git a/docs/source/design.rst b/docs/source/design.rst
@@ -12,6 +12,21 @@ or if you are just looking to access some information that is not currently expo
 .. contents:: Contents
    :local:
 
+.. _initialization:
+
+Initialization
+**************
+
+Use :func:`~kasa.Discover.discover` to perform udp-based broadcast discovery on the network.
+This will return you a list of device instances based on the discovery replies.
+
+If the device's host is already known, you can use to construct a device instance with
+:meth:`~kasa.SmartDevice.connect()`.
+
+When connecting a device with the :meth:`~kasa.SmartDevice.connect()` method, it is recommended to
+pass the device type as well as this allows the library to use the correct device class for the
+device without having to query the device.
+
 .. _update_cycle:
 
 Update Cycle

diff --git a/kasa/cli.py b/kasa/cli.py
@@ -13,14 +13,13 @@
 from kasa import (
     AuthenticationException,
     Credentials,
+    DeviceType,
     Discover,
     SmartBulb,
     SmartDevice,
-    SmartDimmer,
-    SmartLightStrip,
-    SmartPlug,
     SmartStrip,
 )
+from kasa.device_factory import DEVICE_TYPE_TO_CLASS
 
 try:
     from rich import print as _do_echo
@@ -43,13 +42,11 @@ def wrapper(message=None, *args, **kwargs):
 # --json has set it to _nop_echo
 echo = _do_echo
 
-TYPE_TO_CLASS = {
-    "plug": SmartPlug,
-    "bulb": SmartBulb,
-    "dimmer": SmartDimmer,
-    "strip": SmartStrip,
-    "lightstrip": SmartLightStrip,
-}
+DEVICE_TYPES = [
+    device_type.value
+    for device_type in DeviceType
+    if device_type in DEVICE_TYPE_TO_CLASS
+]
 
 click.anyio_backend = "asyncio"
 
@@ -129,7 +126,7 @@ def _device_to_serializable(val: SmartDevice):
     "--type",
     envvar="KASA_TYPE",
     default=None,
-    type=click.Choice(list(TYPE_TO_CLASS), case_sensitive=False),
+    type=click.Choice(DEVICE_TYPES, case_sensitive=False),
 )
 @click.option(
     "--json", default=False, is_flag=True, help="Output raw device response as JSON."
@@ -235,16 +232,19 @@ def _nop_echo(*args, **kwargs):
         return await ctx.invoke(discover, timeout=discovery_timeout)
 
     if type is not None:
-        dev = TYPE_TO_CLASS[type](host, credentials=credentials)
+        device_type = DeviceType.from_value(type)
+        dev = await SmartDevice.connect(
+            host, credentials=credentials, device_type=device_type
+        )
     else:
         echo("No --type defined, discovering..")
         dev = await Discover.discover_single(
             host,
             port=port,
             credentials=credentials,
         )
+        await dev.update()
 
-    await dev.update()
     ctx.obj = dev
 
     if ctx.invoked_subcommand is None:

diff --git a/kasa/device_factory.py b/kasa/device_factory.py
@@ -0,0 +1,121 @@
+"""Device creation by type."""
+
+import logging
+import time
+from typing import Any, Dict, Optional, Type
+
+from .credentials import Credentials
+from .device_type import DeviceType
+from .exceptions import UnsupportedDeviceException
+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,
+}
+
+_LOGGER = logging.getLogger(__name__)
+
+
+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.
+    """
+    debug_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
+
+    if debug_enabled:
+        start_time = time.perf_counter()
+
+    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()
+        if debug_enabled:
+            end_time = time.perf_counter()
+            _LOGGER.debug(
+                "Device %s with known type (%s) took %.2f seconds to connect",
+                host,
+                device_type.value,
+                end_time - start_time,
+            )
+        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()
+    if debug_enabled:
+        end_time = time.perf_counter()
+        _LOGGER.debug(
+            "Device %s with unknown type (%s) took %.2f seconds to connect",
+            host,
+            dev.device_type.value,
+            end_time - start_time,
+        )
+    return dev
+
+
+def get_device_class_from_info(info: Dict[str, Any]) -> 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: Dict[str, Any] = info["system"]["get_sysinfo"]
+    type_: Optional[str] = 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 UnsupportedDeviceException("Unknown device type: %s" % type_)
diff --git a/kasa/device_type.py 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
@@ -20,13 +20,11 @@
 from kasa.json import dumps as json_dumps
 from kasa.json import loads as json_loads
 from kasa.klapprotocol import TPLinkKlap
-from kasa.protocol import TPLinkProtocol, TPLinkSmartHomeProtocol
-from kasa.smartbulb import SmartBulb
+from kasa.protocol import TPLinkSmartHomeProtocol
 from kasa.smartdevice import SmartDevice, SmartDeviceException
-from kasa.smartdimmer import SmartDimmer
-from kasa.smartlightstrip import SmartLightStrip
 from kasa.smartplug import SmartPlug
-from kasa.smartstrip import SmartStrip
+
+from .device_factory import get_device_class_from_info
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -345,78 +343,10 @@ async def discover_single(
         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,
-        protocol_class: Optional[Type[TPLinkProtocol]] = 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 port: Optionally set a different port for the device
-        :param timeout: Timeout for discovery
-        :param credentials: Optionally provide credentials for
-            devices requiring them
-        :param protocol_class: Optionally provide the protocol class
-            to use.
-        :rtype: SmartDevice
-        :return: Object for querying/controlling found device.
-        """
-        unknown_dev = SmartDevice(
-            host=host, port=port, credentials=credentials, timeout=timeout
-        )
-        if protocol_class is not None:
-            unknown_dev.protocol = protocol_class(host, credentials=credentials)
-        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
-        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 UnsupportedDeviceException("Unknown device type: %s" % type_)
+        return get_device_class_from_info(info)
 
     @staticmethod
     def _get_device_instance_legacy(data: bytes, ip: str, port: int) -> SmartDevice: