Repository: mill1000/midea-msmart
Pull request #35 (2023-08-25 20:04:42)

Implement additional type hint and fix discovered issues
Add type hints across the project and address all discovered issues.


Patch:
diff --git a/msmart/base_command.py b/msmart/base_command.py
index 0ba30f6..f3068dc 100644
--- a/msmart/base_command.py
+++ b/msmart/base_command.py
@@ -2,7 +2,7 @@ import logging
 from abc import ABC, abstractmethod
 
 import msmart.crc8 as crc8
-from msmart.const import FrameType
+from msmart.const import DeviceType, FrameType
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -10,12 +10,12 @@ _LOGGER = logging.getLogger(__name__)
 class Command(ABC):
     _message_id = 0
 
-    def __init__(self, device_type=0xAC, frame_type=FrameType.REQUEST):
+    def __init__(self, device_type: DeviceType = DeviceType.AIR_CONDITIONER, frame_type: FrameType = FrameType.REQUEST) -> None:
         self.device_type = device_type
         self.frame_type = frame_type
         self.protocol_version = 0
 
-    def pack(self):
+    def pack(self) -> bytes:
         # Create payload with message id
         payload = self.payload + bytes([self.message_id])
 
@@ -26,7 +26,7 @@ class Command(ABC):
         length = 10 + len(payload_crc)
 
         # Build frame header
-        header = bytearray([
+        header = bytes([
             # Start byte
             0xAA,
             # Length of payload and header
@@ -48,25 +48,25 @@ class Command(ABC):
         ])
 
         # Build frame from header and payload with CRC
-        frame = header + payload_crc
+        frame = bytearray(header + payload_crc)
 
         # Calculate total frame checksum
         frame.append(Command.checksum(frame[1:]))
 
         _LOGGER.debug("Frame data: %s", frame.hex())
 
-        return frame
+        return bytes(frame)
 
-    @staticmethod
-    def checksum(frame):
+    @classmethod
+    def checksum(cls, frame: bytes) -> int:
         return (~sum(frame) + 1) & 0xFF
 
     @property
-    def message_id(self):
+    def message_id(self) -> int:
         Command._message_id += 1
         return Command._message_id & 0xFF
 
     @property
     @abstractmethod
-    def payload(self):
+    def payload(self) -> bytes:
         return bytes()
diff --git a/msmart/base_device.py b/msmart/base_device.py
index 9dafc62..21b62a1 100644
--- a/msmart/base_device.py
+++ b/msmart/base_device.py
@@ -2,22 +2,23 @@ import logging
 import time
 from typing import List, Optional
 
-from msmart.lan import LAN, AuthenticationError, ProtocolError
-from msmart.types import Key, Token
+from msmart.base_command import Command
+from msmart.const import DeviceType
+from msmart.lan import LAN, AuthenticationError, Key, ProtocolError, Token
 
 _LOGGER = logging.getLogger(__name__)
 
 
 class Device():
 
-    def __init__(self, *, ip: str, port: int, device_id: int, **kwargs) -> None:
+    def __init__(self, *, ip: str, port: int, device_id: int, device_type: DeviceType, **kwargs) -> None:
         self._ip = ip
         self._port = port
 
         self._id = device_id
+        self._type = device_type
         self._sn = kwargs.get("sn", None)
         self._name = kwargs.get("name", None)
-        self._type = kwargs.get("type", None)
         self._version = kwargs.get("version", None)
 
         self._lan = LAN(ip, port, device_id)
@@ -39,7 +40,7 @@ class Device():
             _LOGGER.error("Authentication failed. Error: %s", e)
             return False
 
-    async def send_command(self, command: bytes) -> List[object]:
+    async def send_command(self, command: Command) -> Optional[List[bytes]]:
         """Send a command to the device and return any responses."""
 
         data = command.pack()
@@ -92,7 +93,7 @@ class Device():
         return self._lan.key.hex()
 
     @property
-    def type(self) -> Optional[str]:
+    def type(self) -> DeviceType:
         return self._type
 
     @property
diff --git a/msmart/cli.py b/msmart/cli.py
index be1c29b..2ec1f09 100644
--- a/msmart/cli.py
+++ b/msmart/cli.py
@@ -9,7 +9,7 @@ from msmart.discover import Discover
 _LOGGER = logging.getLogger(__name__)
 
 
-async def _discover(ip: str, count: int, account: str, password: str, china: bool, **_kwargs):
+async def _discover(ip: str, count: int, account: str, password: str, china: bool, **_kwargs) -> None:
     """Discover Midea devices and print configuration information."""
 
     _LOGGER.info("msmart version: %s", __version__)
@@ -38,7 +38,7 @@ async def _discover(ip: str, count: int, account: str, password: str, china: boo
         _LOGGER.info("Found device:\n%s", device)
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(description="Discover Midea devices and print device information.",
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument(
diff --git a/msmart/crc8.py b/msmart/crc8.py
index d1eeaaf..9b5c54d 100644
--- a/msmart/crc8.py
+++ b/msmart/crc8.py
@@ -35,7 +35,7 @@ crc8_854_table = [
 ]
 
 
-def calculate(data):
+def calculate(data: bytes) -> int:
     crc_value = 0
     for m in data:
         k = crc_value ^ m
diff --git a/msmart/device/AC/command.py b/msmart/device/AC/command.py
index 29a5638..abf4227 100644
--- a/msmart/device/AC/command.py
+++ b/msmart/device/AC/command.py
@@ -1,12 +1,15 @@
+from __future__ import annotations
+
 import logging
 import math
 import struct
 from collections import namedtuple
 from enum import IntEnum
+from typing import Callable, Union
 
 import msmart.crc8 as crc8
 from msmart.base_command import Command
-from msmart.const import FrameType
+from msmart.const import DeviceType, FrameType
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -52,11 +55,11 @@ class TemperatureType(IntEnum):
 
 
 class GetCapabilitiesCommand(Command):
-    def __init__(self, device_type):
+    def __init__(self, device_type: DeviceType) -> None:
         super().__init__(device_type, frame_type=FrameType.REQUEST)
 
     @property
-    def payload(self):
+    def payload(self) -> bytes:
         return bytes([
             # Get capabilities
             0xB5,
@@ -66,13 +69,13 @@ class GetCapabilitiesCommand(Command):
 
 
 class GetStateCommand(Command):
-    def __init__(self, device_type):
+    def __init__(self, device_type: DeviceType) -> None:
         super().__init__(device_type, frame_type=FrameType.REQUEST)
 
         self.temperature_type = TemperatureType.INDOOR
 
     @property
-    def payload(self):
+    def payload(self) -> bytes:
         return bytes([
             # Get state
             0x41,
@@ -90,7 +93,7 @@ class GetStateCommand(Command):
 
 
 class SetStateCommand(Command):
-    def __init__(self, device_type):
+    def __init__(self, device_type: DeviceType) -> None:
         super().__init__(device_type, frame_type=FrameType.SET)
 
         self.beep_on = True
@@ -106,7 +109,7 @@ class SetStateCommand(Command):
         self.freeze_protection_mode = False
 
     @property
-    def payload(self):
+    def payload(self) -> bytes:
         # Build bepp and power status bytes
         beep = 0x42 if self.beep_on else 0
         power = 0x1 if self.power_on else 0
@@ -164,12 +167,12 @@ class SetStateCommand(Command):
 
 
 class ToggleDisplayCommand(Command):
-    def __init__(self, device_type):
+    def __init__(self, device_type: DeviceType) -> None:
         # For whatever reason, toggle display uses a request type...
         super().__init__(device_type, frame_type=FrameType.REQUEST)
 
     @property
-    def payload(self):
+    def payload(self) -> bytes:
         # Payload taken directly from dudanov/MideaUART
         return bytes([
             # Get state
@@ -184,21 +187,21 @@ class ToggleDisplayCommand(Command):
 
 
 class Response():
-    def __init__(self, payload: memoryview):
+    def __init__(self, payload: memoryview) -> None:
         # Set ID and copy the payload
         self._id = payload[0]
         self._payload = bytes(payload)
 
     @property
-    def id(self):
+    def id(self) -> int:
         return self._id
 
     @property
-    def payload(self):
+    def payload(self) -> bytes:
         return self._payload
 
-    @staticmethod
-    def validate(frame: memoryview):
+    @classmethod
+    def validate(cls, frame: memoryview) -> None:
         # Validate frame checksum
         frame_checksum = Command.checksum(frame[1:-1])
         if frame_checksum != frame[-1]:
@@ -216,8 +219,8 @@ class Response():
             raise InvalidResponseException(
                 f"Payload '{payload.hex()}' failed CRC and checksum. Received: 0x{payload[-1]:X}, Expected: 0x{payload_crc:X} or 0x{payload_checksum:X}.")
 
-    @staticmethod
-    def construct(frame: bytes):
+    @classmethod
+    def construct(cls, frame: bytes) -> Union[StateResponse, CapabilitiesResponse, Response]:
         # Build a memoryview of the frame for zero-copy slicing
         with memoryview(frame) as frame_mv:
             # Ensure frame is valid before parsing
@@ -235,7 +238,7 @@ class Response():
 
 
 class CapabilitiesResponse(Response):
-    def __init__(self, payload: memoryview):
+    def __init__(self, payload: memoryview) -> None:
         super().__init__(payload)
 
         self._capabilities = {}
@@ -246,14 +249,14 @@ class CapabilitiesResponse(Response):
 
         _LOGGER.debug("Supported capabilities: %s", self._capabilities)
 
-    def _parse_capabilities(self, payload: memoryview):
+    def _parse_capabilities(self, payload: memoryview) -> None:
         # Clear existing capabilities
         self._capabilities.clear()
 
         # Define some local functions to parse capability values
-        def get_bool(v): return v != 0
-        def get_value(w): return lambda v: v == w
-        def get_no_value(w): return lambda v: v != w
+        def get_bool(v) -> bool: return v != 0
+        def get_value(w) -> Callable[[int], bool]: return lambda v: v == w
+        def get_no_value(w) -> Callable[[int], bool]: return lambda v: v != w
 
         # Define a named tuple that represents a decoder
         reader = namedtuple("decoder", "name read")
@@ -375,62 +378,62 @@ class CapabilitiesResponse(Response):
             caps = caps[3+size:]
 
     @property
-    def swing_horizontal(self):
+    def swing_horizontal(self) -> bool:
         return self._capabilities.get("swing_horizontal", False)
 
     @property
-    def swing_vertical(self):
+    def swing_vertical(self) -> bool:
         return self._capabilities.get("swing_vertical", False)
 
     @property
-    def swing_both(self):
+    def swing_both(self) -> bool:
         return self.swing_vertical and self.swing_horizontal
 
     @property
-    def dry_mode(self):
+    def dry_mode(self) -> bool:
         return self._capabilities.get("dry_mode", False)
 
     @property
-    def cool_mode(self):
+    def cool_mode(self) -> bool:
         return self._capabilities.get("cool_mode", False)
 
     @property
-    def heat_mode(self):
+    def heat_mode(self) -> bool:
         return self._capabilities.get("heat_mode", False)
 
     @property
-    def auto_mode(self):
+    def auto_mode(self) -> bool:
         return self._capabilities.get("auto_mode", False)
 
     @property
-    def eco_mode(self):
+    def eco_mode(self) -> bool:
         return self._capabilities.get("eco_mode", False) or self._capabilities.get("eco_mode_2", False)
 
     @property
-    def turbo_mode(self):
+    def turbo_mode(self) -> bool:
         return self._capabilities.get("turbo_heat", False) or self._capabilities.get("turbo_cool", False)
 
     @property
-    def display_control(self):
+    def display_control(self) -> bool:
         return self._capabilities.get("light_control", False)
 
     @property
-    def min_temperature(self):
+    def min_temperature(self) -> int:
         mode = ["cool", "auto", "heat"]
         return min([self._capabilities.get(f"{m}_min_temperature", 16) for m in mode])
 
     @property
-    def max_temperature(self):
+    def max_temperature(self) -> int:
         mode = ["cool", "auto", "heat"]
         return max([self._capabilities.get(f"{m}_max_temperature", 30) for m in mode])
 
     @property
-    def freeze_protection_mode(self):
+    def freeze_protection_mode(self) -> bool:
         return self._capabilities.get("freeze_protection", False)
 
 
 class StateResponse(Response):
-    def __init__(self, payload: memoryview):
+    def __init__(self, payload: memoryview) -> None:
         super().__init__(payload)
 
         self.power_on = None
@@ -452,7 +455,7 @@ class StateResponse(Response):
 
         self._parse(payload)
 
-    def _parse(self, payload: memoryview):
+    def _parse(self, payload: memoryview) -> None:
 
         self.power_on = bool(payload[1] & 0x1)
         # self.imode_resume = payload[1] & 0x4
diff --git a/msmart/device/AC/device.py b/msmart/device/AC/device.py
index 77113d6..98298dd 100644
--- a/msmart/device/AC/device.py
+++ b/msmart/device/AC/device.py
@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import logging
 from enum import IntEnum
-from typing import Any, List, Optional
+from typing import Any, List, Optional, cast
 
 from msmart.base_device import Device
 from msmart.const import DeviceType
@@ -15,26 +17,26 @@ _LOGGER = logging.getLogger(__name__)
 
 
 class IntEnumHelper(IntEnum):
-    """Helper class to convert IntEnums to strings."""
+    """Helper class to convert IntEnums to/from strings."""
     @classmethod
-    def list(cls) -> List[IntEnum]:
+    def list(cls) -> List[IntEnumHelper]:
         return list(map(lambda c: c, cls))
 
     @classmethod
-    def get_from_value(cls, value, default=None) -> Any:
+    def get_from_value(cls, value: Optional[int], default: IntEnumHelper) -> IntEnumHelper:
         try:
-            return cls(value)
+            return cls(cast(int, value))
         except ValueError:
             _LOGGER.debug("Unknown %s: %d", cls, value)
-            return cls(default)
+            return default
 
     @classmethod
-    def get_from_name(cls, name, default=None) -> Any:
+    def get_from_name(cls, name: str, default: IntEnumHelper) -> IntEnumHelper:
         try:
             return cls[name]
         except KeyError:
             _LOGGER.debug("Unknown %s: %d", cls, name)
-            return cls[default]
+            return default
 
 
 class AirConditioner(Device):
@@ -48,12 +50,12 @@ class AirConditioner(Device):
         SILENT = 20
 
         @classmethod
-        def get_from_value(cls, value, default=AUTO) -> IntEnum:
-            return super().get_from_value(value, default)
+        def get_from_value(cls, value: Optional[int], default=AUTO) -> AirConditioner.FanSpeed:
+            return cast(cls, super().get_from_value(value, default))
 
         @classmethod
-        def get_from_name(cls, name, default=AUTO) -> IntEnum:
-            return super().get_from_name(name, default)
+        def get_from_name(cls, name: str, default=AUTO) -> AirConditioner.FanSpeed:
+            return cast(cls, super().get_from_name(name, default))
 
     class OperationalMode(IntEnumHelper):
         AUTO = 1
@@ -63,12 +65,12 @@ class AirConditioner(Device):
         FAN_ONLY = 5
 
         @classmethod
-        def get_from_value(cls, value, default=FAN_ONLY) -> IntEnum:
-            return super().get_from_value(value, default)
+        def get_from_value(cls, value: Optional[int], default=FAN_ONLY) -> AirConditioner.OperationalMode:
+            return cast(cls, super().get_from_value(value, default))
 
         @classmethod
-        def get_from_name(cls, name, default=FAN_ONLY) -> IntEnum:
-            return super().get_from_name(name, default)
+        def get_from_name(cls, name: str, default=FAN_ONLY) -> AirConditioner.OperationalMode:
+            return cast(cls, super().get_from_name(name, default))
 
     class SwingMode(IntEnumHelper):
         OFF = 0x0
@@ -77,18 +79,19 @@ class AirConditioner(Device):
         BOTH = 0xF
 
         @classmethod
-        def get_from_value(cls, value, default=OFF) -> IntEnum:
-            return super().get_from_value(value, default)
+        def get_from_value(cls, value: Optional[int], default=OFF) -> AirConditioner.SwingMode:
+            return cast(cls, super().get_from_value(value, default))
 
         @classmethod
-        def get_from_name(cls, name, default=OFF) -> IntEnum:
-            return super().get_from_name(name, default)
+        def get_from_name(cls, name: str, default=OFF) -> AirConditioner.SwingMode:
+            return cast(cls, super().get_from_name(name, default))
 
     def __init__(self, ip: str, device_id: int,  port: int, **kwargs) -> None:
-        # Ensure type is set
-        kwargs["type"] = DeviceType.AIR_CONDITIONER.value
+        # Remove possible duplicate device_type kwarg
+        kwargs.pop("device_type", None)
 
-        super().__init__(ip=ip, port=port, device_id=device_id, **kwargs)
+        super().__init__(ip=ip, port=port, device_id=device_id,
+                         device_type=DeviceType.AIR_CONDITIONER, **kwargs)
 
         self._updating = False
         self._keep_last_known_online_state = False
@@ -109,8 +112,10 @@ class AirConditioner(Device):
         self._filter_alert = False
 
         # Support all known modes initially
-        self._supported_op_modes = AirConditioner.OperationalMode.list()
-        self._supported_swing_modes = AirConditioner.SwingMode.list()
+        self._supported_op_modes = cast(
+            List[AirConditioner.OperationalMode], AirConditioner.OperationalMode.list())
+        self._supported_swing_modes = cast(
+            List[AirConditioner.SwingMode], AirConditioner.SwingMode.list())
         self._supports_eco = True
         self._supports_turbo = True
         self._supports_freeze_protection_mode = True
@@ -165,8 +170,10 @@ class AirConditioner(Device):
             self._supported = True
 
             if response.id == ResponseId.STATE:
+                response = cast(StateResponse, response)
                 self.update(response)
             elif response.id == ResponseId.CAPABILITIES:
+                response = cast(CapabilitiesResponse, response)
                 self.update_capabilities(response)
             else:
                 _LOGGER.debug("Ignored unknown response from %s:%d: %s",
@@ -195,18 +202,24 @@ class AirConditioner(Device):
             if self._freeze_protection_mode and not self._supports_freeze_protection_mode:
                 _LOGGER.warning("Device is not capable of freeze protection.")
 
+            # Define function to return value or a default if value is None
+            def or_default(v, d) -> Any: return v if v is not None else d
+
             cmd = SetStateCommand(self.type)
             cmd.beep_on = self._beep_on
-            cmd.power_on = self._power_state
-            cmd.target_temperature = self._target_temperature
+            cmd.power_on = or_default(self._power_state, False)
+            cmd.target_temperature = or_default(
+                self._target_temperature, 25)  # TODO?
             cmd.operational_mode = self._operational_mode
             cmd.fan_speed = self._fan_speed
             cmd.swing_mode = self._swing_mode
-            cmd.eco_mode = self._eco_mode
-            cmd.turbo_mode = self._turbo_mode
-            cmd.freeze_protection_mode = self._freeze_protection_mode
-            cmd.sleep_mode = self._sleep_mode
-            cmd.fahrenheit = self._fahrenheit_unit
+            cmd.eco_mode = or_default(self._eco_mode, False)
+            cmd.turbo_mode = or_default(self._turbo_mode, False)
+            cmd.freeze_protection_mode = or_default(
+                self._freeze_protection_mode, False)
+            cmd.sleep_mode = or_default(self._sleep_mode, False)
+            cmd.fahrenheit = or_default(self._fahrenheit_unit, False)
+
             await self.send_command(cmd, self._defer_update)
         finally:
             self._updating = False
diff --git a/msmart/discover.py b/msmart/discover.py
index cb54378..40195f5 100644
--- a/msmart/discover.py
+++ b/msmart/discover.py
@@ -4,9 +4,9 @@ import ipaddress
 import logging
 import socket
 import xml.etree.ElementTree as ET
-from typing import Any, Dict, Optional, Type, cast
+from typing import Any, Dict, List, Optional, Type, cast
 
-from msmart.cloud import ApiError, Cloud, CloudError
+from msmart.cloud import Cloud, CloudError
 from msmart.const import (DEVICE_INFO_MSG, DISCOVERY_MSG,
                           OPEN_MIDEA_APP_ACCOUNT, OPEN_MIDEA_APP_PASSWORD,
                           DeviceType)
@@ -18,16 +18,22 @@ _LOGGER = logging.getLogger(__name__)
 _IPV4_BROADCAST = "255.255.255.255"
 
 
+class DiscoverError(Exception):
+    pass
+
+
 class _V1DeviceInfoProtocol(asyncio.Protocol):
     """V1 device info protocol."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._transport = None
         self.response = None
 
     def connection_made(self, transport) -> None:
         """Send device info request on connection."""
 
+        transport = cast(asyncio.Transport, transport)
+
         self._transport = transport
         transport.write(DEVICE_INFO_MSG)
 
@@ -51,7 +57,7 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
         target: str = _IPV4_BROADCAST,
         discovery_packets: int = 3,
         interface: Optional[str] = None,
-    ):
+    ) -> None:
         self._transport = None
         self._discovery_packets = discovery_packets
         self._interface = interface
@@ -63,6 +69,8 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
     def connection_made(self, transport) -> None:
         """Set socket options for broadcasting."""
 
+        transport = cast(asyncio.DatagramTransport, transport)
+
         self._transport = transport
 
         # Set broadcast if broadcasting
@@ -71,6 +79,7 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
             sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
 
         if self._interface is not None:
+            sock = transport.get_extra_info("socket")
             sock.setsockopt(
                 socket.SOL_SOCKET, socket.SO_BINDTODEVICE, self._interface.encode()
             )
@@ -80,6 +89,9 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
     def _send_discovery(self) -> None:
         """Send discovery messages to the target."""
 
+        # Transport should always exist
+        assert self._transport is not None
+
         for port in [6445, 20086]:
             _LOGGER.debug("Discovery sent to %s:%d.", self._target, port)
             for _ in range(self._discovery_packets):
@@ -101,7 +113,7 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
         try:
             # pylint: disable=protected-access
             version = Discover._get_device_version(data)
-        except Discover.UnknownDeviceVersion:
+        except DiscoverError:
             _LOGGER.error("Unknown device version for %s.", ip)
             return
 
@@ -123,9 +135,6 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
 class Discover:
     """Discover Midea smart devices on the local network."""
 
-    class UnknownDeviceVersion(Exception):
-        """Exception for unknown device version."""
-
     _account = OPEN_MIDEA_APP_ACCOUNT
     _password = OPEN_MIDEA_APP_PASSWORD
     _lock = asyncio.Lock()
@@ -143,7 +152,7 @@ class Discover:
         account=None,
         password=None,
         auto_connect=True
-    ) -> [Device]:
+    ) -> List[Device]:
         """Discover devices via broadcast."""
 
         # Always use a new cloud connection
@@ -188,7 +197,7 @@ class Discover:
         cls,
         host,
         **kwargs
-    ) -> Device:
+    ) -> Optional[Device]:
         """Discover a single device by hostname or IP."""
 
         devices = await Discover.discover(target=host, **kwargs)
@@ -210,7 +219,7 @@ class Discover:
             raise ValueError("Both account and password must be specified.")
 
     @classmethod
-    async def _get_cloud(cls) -> Cloud:
+    async def _get_cloud(cls) -> Optional[Cloud]:
         """Return a cloud connection, creating it if necessary."""
 
         async with cls._lock:
@@ -244,7 +253,7 @@ class Discover:
             elif start_of_packet == b'\x83\x70':
                 return 3
 
-        raise Discover.UnknownDeviceVersion()
+        raise DiscoverError()
 
     @classmethod
     async def _get_device_info(cls, ip: str, version: int, data: bytes) -> Dict[str, Any]:
@@ -259,7 +268,11 @@ class Discover:
             with memoryview(data) as data_mv:
                 root = ET.fromstring(data_mv)
 
-            port = root.find("body/device").attrib["port"]
+            device = root.find("body/device")
+            if device is None:
+                raise DiscoverError("Could not find 'body/device' in XML.")
+
+            port = int(device.attrib["port"])
 
             loop = asyncio.get_event_loop()
             transport, protocol = await loop.create_connection(
@@ -273,9 +286,13 @@ class Discover:
             finally:
                 transport.close()
 
+            if protocol.response is None:
+                raise DiscoverError(
+                    f"No device info response from {ip}:{port}.")
+
             # Parse response
-            root = ET.fromstring(protocol.response)
-            _LOGGER.debug("Device info response:\n%s", ET.indent(root))
+            root = ET.fromstring(protocol.response.decode())
+            _LOGGER.debug("Device info response:\n%s", ET.tostring(root))
 
             raise NotImplementedError("V1 device not supported yet.")
 
@@ -295,9 +312,9 @@ class Discover:
                 # Attempt to decrypt the packet
                 try:
                     decrypted_data = Security.decrypt_aes(encrypted_data)
-                except ValueError:
-                    _LOGGER.error("Discovery packet decrypt failed.")
-                    return
+                except ValueError as e:
+                    raise DiscoverError(
+                        "Failed to decrypt discovery resonse.") from e
 
             with memoryview(decrypted_data) as decrypted_mv:
                 _LOGGER.debug("Decrypted data from %s: %s",
@@ -322,7 +339,7 @@ class Discover:
                 device_type = int(name.split('_')[1], 16)
 
             # Return dictionary of device info
-            return {"ip": ip_address, "port": port, "device_id": device_id, "name": name, "sn": sn, "type": device_type, "version": version}
+            return {"ip": ip_address, "port": port, "device_id": device_id, "name": name, "sn": sn, "device_type": device_type, "version": version}
 
     @classmethod
     def _get_device_class(cls, device_type: int) -> Type[Device]:
@@ -346,7 +363,8 @@ class Discover:
 
         # Try authenticating with udpids generated from both endians
         for endian in ["little", "big"]:
-            udpid = Security.udpid(dev.id.to_bytes(6, endian)).hex()
+            udpid = Security.udpid(dev.id.to_bytes(
+                6, endian)).hex()  # type: ignore
 
             _LOGGER.debug(
                 "Fetching token and key for udpid '%s' (%s).", udpid, endian)
@@ -362,18 +380,18 @@ class Discover:
         return False
 
     @classmethod
-    async def _get_device(cls, ip: str, version: int, data: bytes) -> Device:
+    async def _get_device(cls, ip: str, version: int, data: bytes) -> Optional[Device]:
         """Get device information and construct a device instance from the discovery response data."""
 
         # Fetch device information
         try:
             info = await Discover._get_device_info(ip, version, data)
-        except NotImplementedError as e:
+        except (DiscoverError, NotImplementedError) as e:
             _LOGGER.error(e)
             return None
 
         # Get device class corresponding to type
-        device_class = Discover._get_device_class(info["type"])
+        device_class = Discover._get_device_class(info["device_type"])
 
         # Build device, authenticate as needed and refresh
         dev = device_class(**info)
diff --git a/msmart/lan.py b/msmart/lan.py
index f4ed2dc..1c77a5d 100644
--- a/msmart/lan.py
+++ b/msmart/lan.py
@@ -5,14 +5,15 @@ import struct
 from datetime import datetime
 from enum import IntEnum
 from hashlib import md5, sha256
-from typing import cast
+from typing import List, Optional, Union, cast
 
 from Crypto.Cipher import AES
 from Crypto.Random import get_random_bytes
 from Crypto.Util import Padding
 from Crypto.Util.strxor import strxor
 
-from msmart.types import Key, Token
+Token = Union[str, bytes, None]
+Key = Union[str, bytes, None]
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -47,18 +48,18 @@ class _LanProtocol(asyncio.Protocol):
     #   16 byte MD5 of packet + fixed key
     #
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._transport = None
 
         self._peer = None
         self._queue = asyncio.Queue()
 
     @property
-    def peer(self):
+    def peer(self) -> Optional[str]:
         return self._peer
 
     @property
-    def alive(self):
+    def alive(self) -> bool:
         if self._transport is None or self._transport.is_closing():
             return False
 
@@ -72,6 +73,8 @@ class _LanProtocol(asyncio.Protocol):
     def connection_made(self, transport) -> None:
         """Handle connection events."""
 
+        transport = cast(asyncio.Transport, transport)
+
         # Save transport for later
         self._transport = transport
 
@@ -157,7 +160,7 @@ class _LanProtocolV3(_LanProtocol):
     #    but not in the size field
     #  - When used for device command/response, the payload contains a V2 packet
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
         self._packet_id = 0
@@ -211,9 +214,12 @@ class _LanProtocolV3(_LanProtocol):
                 # Queue the received packet
                 self._queue.put_nowait(packet.tobytes())
 
-    def _decode_encrypted_response(self, packet: memoryview):
+    def _decode_encrypted_response(self, packet: memoryview) -> bytes:
         """Decode an encrypted response packet."""
 
+        # We should always have a key by the time we're received data
+        assert self._local_key is not None
+
         # Extract header, encrypted payload and hash
         header = packet[:6]
         payload = packet[6:-32]
@@ -236,9 +242,9 @@ class _LanProtocolV3(_LanProtocol):
             pad = header[5] >> 4
 
             # Get the frame from payload
-            return payload[2:-pad]
+            return payload[2:-pad].tobytes()
 
-    def _decode_handshake_response(self, packet: memoryview):
+    def _decode_handshake_response(self, packet: memoryview) -> bytes:
         """Decode a handshake response packet."""
 
         # Get payload from packet
@@ -247,7 +253,7 @@ class _LanProtocolV3(_LanProtocol):
         # Return remaining raw payload
         return payload[2:].tobytes()
 
-    def _process_packet(self, packet: memoryview):
+    def _process_packet(self, packet: memoryview) -> bytes:
         """Process a received packet based on its type."""
 
         if packet[:2] != b"\x83\x70":
@@ -278,7 +284,7 @@ class _LanProtocolV3(_LanProtocol):
         with memoryview(packet) as packet_mv:
             return self._process_packet(packet_mv)
 
-    def _build_header(self, length: int, extra: bytes):
+    def _build_header(self, length: int, extra: bytes) -> bytes:
         # Build header
         header = b"\x83\x70"
         header += length.to_bytes(2, "big")
@@ -287,9 +293,12 @@ class _LanProtocolV3(_LanProtocol):
 
         return header
 
-    def _encode_encrypted_request(self, packet_id: int, data: bytes):
+    def _encode_encrypted_request(self, packet_id: int, data: bytes) -> bytes:
         """Encode an encrypted request packet."""
 
+        # We should always have a key before sending data
+        assert self._local_key is not None
+
         # Compute required padding for 16 byte alignment
         # Include 2 bytes for packet ID in total length
         remainder = (len(data) + 2) % 16
@@ -308,7 +317,7 @@ class _LanProtocolV3(_LanProtocol):
         calc_hash = sha256(header + payload).digest()
         return header + Security.encrypt_aes_cbc(self._local_key, payload) + calc_hash
 
-    def _encode_handshake_request(self, packet_id: int, data: bytes):
+    def _encode_handshake_request(self, packet_id: int, data: bytes) -> bytes:
         """Encode a handshake request packet."""
 
         # Build header
@@ -342,7 +351,7 @@ class _LanProtocolV3(_LanProtocol):
         self._packet_id += 1
         self._packet_id &= 0xFFF  # Mask to 12 bits
 
-    def _get_local_key(self, key: Key, data: memoryview):
+    def _get_local_key(self, key: bytes, data: memoryview) -> bytes:
 
         if len(data) != 64:
             raise AuthenticationError(
@@ -362,7 +371,7 @@ class _LanProtocolV3(_LanProtocol):
         # Construct the local key
         return strxor(decrypted_payload, key)
 
-    async def authenticate(self, token: Token, key: Key):
+    async def authenticate(self, token: Optional[bytes], key: Optional[bytes]) -> None:
 
         # Raise an exception if trying to auth without any token or key
         if not token or not key:
@@ -386,7 +395,7 @@ class _LanProtocolV3(_LanProtocol):
 class LAN:
     RETRIES = 3
 
-    def __init__(self, ip: str, port: int, device_id: int):
+    def __init__(self, ip: str, port: int, device_id: int) -> None:
         self._ip = ip
         self._port = port
         self._device_id = device_id
@@ -397,14 +406,14 @@ class LAN:
         self._protocol = None
 
     @property
-    def token(self) -> bytes:
+    def token(self) -> Optional[bytes]:
         return self._token
 
     @property
-    def key(self) -> bytes:
+    def key(self) -> Optional[bytes]:
         return self._key
 
-    async def _connect(self):
+    async def _connect(self) -> None:
         _LOGGER.info("Creating new connection to %s:%s", self._ip, self._port)
 
         protocol_class = _LanProtocolV3 if self._protocol_version == 3 else _LanProtocol
@@ -419,12 +428,12 @@ class LAN:
         else:
             self._protocol = cast(_LanProtocol, protocol)
 
-    def _disconnect(self):
+    def _disconnect(self) -> None:
         if self._protocol:
             self._protocol.disconnect()
             self._protocol = None
 
-    async def authenticate(self, token: Token = None, key: Key = None, retries: int = RETRIES):
+    async def authenticate(self, token: Token = None, key: Key = None, retries: int = RETRIES) -> None:
         """Authenticate against a V3 device. Use cached token and key unless provided a new token and key."""
 
         # Use existing token and key if none provided
@@ -433,7 +442,7 @@ class LAN:
             key = self._key
         else:
             # Define a lambda to convert hex strings to bytes
-            def convert(x):
+            def convert(x: Union[Token, Key]) -> Optional[bytes]:
                 return bytes.fromhex(x) if isinstance(x, str) else x
 
             # Ensure passed token and key are in byte form
@@ -449,12 +458,15 @@ class LAN:
             self._protocol_version = 3
             await self._connect()
 
+        # A protocl should exist at this point
+        assert self._protocol is not None
+
         _LOGGER.info("Authenticating with %s.", self._protocol.peer)
 
         # Attempt to authenticate
         while retries > 0:
             try:
-                await self._protocol.authenticate(token, key)
+                await cast(_LanProtocolV3, self._protocol).authenticate(token, key)
                 break
             except (TimeoutError, asyncio.TimeoutError) as e:
                 if retries > 1:
@@ -470,9 +482,12 @@ class LAN:
         # Sleep briefly before requesting more data
         await asyncio.sleep(1)
 
-    async def _read(self, **kwargs):
+    async def _read(self, **kwargs) -> bytes:
         """Read and decode a frame from the protocol."""
 
+        # A protocl should exist at this point
+        assert self._protocol is not None
+
         # Await a response
         packet = await self._protocol.read(**kwargs)
         _LOGGER.debug("Received packet from %s: %s",
@@ -485,7 +500,7 @@ class LAN:
 
         return response
 
-    async def send(self, data: bytes, retries: int = RETRIES):
+    async def send(self, data: bytes, retries: int = RETRIES) -> List[bytes]:
         """Send data via the LAN protocol. Connecting to the peer if necessary."""
 
         # Connect if protocol doesn't exist or is dead
@@ -497,6 +512,9 @@ class LAN:
             if self._protocol_version == 3:
                 await self.authenticate()
 
+        # A protocl should exist at this point
+        assert self._protocol is not None
+
         # Encode frame to packet
         packet = _Packet.encode(self._device_id, data)
 
@@ -534,33 +552,33 @@ class Security:
     ENC_KEY = md5(SIGN_KEY).digest()
 
     @classmethod
-    def decrypt_aes_cbc(cls, key: bytes, data: bytes):
+    def decrypt_aes_cbc(cls, key: bytes, data: bytes) -> bytes:
         return AES.new(key, AES.MODE_CBC, iv=bytes(16)).decrypt(data)
 
     @classmethod
-    def encrypt_aes_cbc(cls, key: bytes, data: bytes):
+    def encrypt_aes_cbc(cls, key: bytes, data: bytes) -> bytes:
         return AES.new(key, AES.MODE_CBC, iv=bytes(16)).encrypt(data)
 
     @classmethod
-    def decrypt_aes(cls, data: bytes):
+    def decrypt_aes(cls, data: bytes) -> bytes:
         cipher = AES.new(Security.ENC_KEY, AES.MODE_ECB)
 
         # Decrypt and remove padding
         return Padding.unpad(cipher.decrypt(data), 16)
 
     @classmethod
-    def encrypt_aes(cls, data: bytes):
+    def encrypt_aes(cls, data: bytes) -> bytes:
         cipher = AES.new(Security.ENC_KEY, AES.MODE_ECB)
 
         # Encrypt the padded data
         return cipher.encrypt(Padding.pad(data, 16))
 
     @classmethod
-    def sign(cls, data: bytes):
+    def sign(cls, data: bytes) -> bytes:
         return md5(data + Security.SIGN_KEY).digest()
 
     @classmethod
-    def udpid(cls, device_id: bytes):
+    def udpid(cls, device_id: bytes) -> bytes:
         with memoryview(sha256(device_id).digest()) as mv_hash:
             return strxor(mv_hash[:16], mv_hash[16:])
 
@@ -569,7 +587,7 @@ class _Packet:
     """Class to encode/decode command frames to packets."""
 
     @classmethod
-    def encode(cls, device_id: int, command: bytes):
+    def encode(cls, device_id: int, command: bytes) -> bytes:
         """Encode a command frame to a LAN packet."""
         # Encrypt command
         encrypted_payload = Security.encrypt_aes(command)
@@ -592,7 +610,7 @@ class _Packet:
         return packet + Security.sign(packet)
 
     @classmethod
-    def decode(cls, data: bytes):
+    def decode(cls, data: bytes) -> bytes:
         """Decode a LAN packet to a command frame."""
 
         with memoryview(data) as packet:
@@ -622,7 +640,7 @@ class _Packet:
             return Security.decrypt_aes(encrypted_frame)
 
     @classmethod
-    def _timestamp(cls):
+    def _timestamp(cls) -> bytes:
         now = datetime.utcnow()
 
         # Each byte is a 2 digit component of the timestamp
diff --git a/msmart/types.py b/msmart/types.py
deleted file mode 100644
index 2daa2c1..0000000
--- a/msmart/types.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from typing import Union
-
-Token = Union[str, bytes, None]
-Key = Union[str, bytes, None]


Test patch:
diff --git a/msmart/device/AC/test_command.py b/msmart/device/AC/test_command.py
index c1ebb31..b42e818 100644
--- a/msmart/device/AC/test_command.py
+++ b/msmart/device/AC/test_command.py
@@ -1,21 +1,22 @@
 import unittest
+from typing import Union, cast
 
-from .command import Response
+from .command import CapabilitiesResponse, Response, StateResponse
 
 
 class TestStateResponse(unittest.TestCase):
-    def assertHasAttr(self, obj, attr):
+    def assertHasAttr(self, obj, attr) -> None:
         """Assert that the supplied object has the supplied attribute."""
         self.assertTrue(hasattr(obj, attr),
                         msg=f"Object {obj} lacks attribute '{attr}'.")
 
-    def _test_build_response(self, msg):
+    def _test_build_response(self, msg) -> Union[StateResponse, CapabilitiesResponse, Response]:
         resp = Response.construct(msg)
         self.assertIsNotNone(resp)
 
         return resp
 
-    def _test_check_attributes(self, resp):
+    def _test_check_attributes(self, resp) -> None:
         expected_attrs = ["power_on", "target_temperature", "operational_mode",
                           "fan_speed", "swing_mode", "turbo_mode", "eco_mode",
                           "sleep_mode", "fahrenheit", "indoor_temperature",
@@ -24,44 +25,62 @@ class TestStateResponse(unittest.TestCase):
         for attr in expected_attrs:
             self.assertHasAttr(resp, attr)
 
-    def _test_response(self, msg):
+    def _test_response(self, msg) -> Union[StateResponse, CapabilitiesResponse, Response]:
         resp = self._test_build_response(msg)
         self._test_check_attributes(resp)
         return resp
 
-    def test_message_checksum(self):
+    def test_message_checksum(self) -> None:
         # V3 state response with checksum as CRC, and shorter than expected
         TEST_MESSAGE_CHECKSUM_AS_CRC = bytes.fromhex(
             "aa1eac00000000000003c0004b1e7f7f000000000069630000000000000d33")
 
         resp = self._test_response(TEST_MESSAGE_CHECKSUM_AS_CRC)
 
+        # Assert response is a state response
+        self.assertEqual(type(resp), StateResponse)
+
+        # Supress type errors
+        resp = cast(StateResponse, resp)
+
         self.assertEqual(resp.target_temperature, 27.0)
         self.assertEqual(resp.indoor_temperature, 27.5)
         self.assertEqual(resp.outdoor_temperature, 24.5)
 
-    def test_message_v2(self):
+    def test_message_v2(self) -> None:
         # V2 state response
         TEST_MESSAGE_V2 = bytes.fromhex(
             "aa22ac00000000000303c0014566000000300010045eff00000000000000000069fdb9")
 
         resp = self._test_response(TEST_MESSAGE_V2)
 
+        # Assert response is a state response
+        self.assertEqual(type(resp), StateResponse)
+
+        # Supress type errors
+        resp = cast(StateResponse, resp)
+
         self.assertEqual(resp.target_temperature, 21.0)
         self.assertEqual(resp.indoor_temperature, 22.0)
         self.assertEqual(resp.outdoor_temperature, None)
 
-    def test_message_v3(self):
+    def test_message_v3(self) -> None:
         # V3 state response
         TEST_MESSAGE_V3 = bytes.fromhex(
             "aa23ac00000000000303c00145660000003c0010045c6b20000000000000000000020d79")
 
         resp = self._test_response(TEST_MESSAGE_V3)
 
+        # Assert response is a state response
+        self.assertEqual(type(resp), StateResponse)
+
+        # Supress type errors
+        resp = cast(StateResponse, resp)
+
         self.assertEqual(resp.target_temperature, 21.0)
         self.assertEqual(resp.indoor_temperature, 21.0)
         self.assertEqual(resp.outdoor_temperature, 28.5)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/msmart/tests/test_discover.py b/msmart/tests/test_discover.py
index 6db5d10..0bd5a3a 100644
--- a/msmart/tests/test_discover.py
+++ b/msmart/tests/test_discover.py
@@ -20,19 +20,22 @@ class TestDiscover(unittest.IsolatedAsyncioTestCase):
 
         # Check info matches
         info = await Discover._get_device_info(IP_ADDRESS, version, DISCOVER_RESPONSE_V2)
-
         self.assertIsNotNone(info)
+
+        # Stop type errors
+        assert info is not None
+
         self.assertEqual(info["ip"], IP_ADDRESS)
         self.assertEqual(info["port"], 6444)
 
         self.assertEqual(info["device_id"], 15393162840672)
-        self.assertEqual(info["type"], DeviceType.AIR_CONDITIONER)
+        self.assertEqual(info["device_type"], DeviceType.AIR_CONDITIONER)
 
         self.assertEqual(info["name"], "net_ac_F7B4")
         self.assertEqual(info["sn"], "000000P0000000Q1F0C9D153F7B40000")
 
         # Check class is correct
-        device_class = Discover._get_device_class(info["type"])
+        device_class = Discover._get_device_class(info["device_type"])
         self.assertEqual(device_class, AC)
 
         # Check that device can be built
@@ -51,19 +54,22 @@ class TestDiscover(unittest.IsolatedAsyncioTestCase):
 
         # Check info matches
         info = await Discover._get_device_info(IP_ADDRESS, version, DISCOVER_RESPONSE_V3)
-
         self.assertIsNotNone(info)
+
+        # Stop type errors
+        assert info is not None
+
         self.assertEqual(info["ip"], IP_ADDRESS)
         self.assertEqual(info["port"], 6444)
 
         self.assertEqual(info["device_id"], 147334558165565)
-        self.assertEqual(info["type"], DeviceType.AIR_CONDITIONER)
+        self.assertEqual(info["device_type"], DeviceType.AIR_CONDITIONER)
 
         self.assertEqual(info["name"], "net_ac_63BA")
         self.assertEqual(info["sn"], "000000P0000000Q1B88C29C963BA0000")
 
         # Check class is correct
-        device_class = Discover._get_device_class(info["type"])
+        device_class = Discover._get_device_class(info["device_type"])
         self.assertEqual(device_class, AC)
 
         # Check that device can be built
@@ -71,5 +77,5 @@ class TestDiscover(unittest.IsolatedAsyncioTestCase):
         self.assertIsNotNone(device)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/msmart/tests/test_lan.py b/msmart/tests/test_lan.py
index bff8ad4..ade32ee 100644
--- a/msmart/tests/test_lan.py
+++ b/msmart/tests/test_lan.py
@@ -43,7 +43,8 @@ class TestEncodeDecode(unittest.IsolatedAsyncioTestCase):
         protocol = _LanProtocolV3()
         protocol._local_key = LOCAL_KEY
 
-        payload = protocol._process_packet(PACKET)
+        with memoryview(PACKET) as mv_packet:
+            payload = protocol._process_packet(mv_packet)
         self.assertIsNotNone(payload)
         self.assertEqual(payload, EXPECTED_PAYLOAD)
 
@@ -52,5 +53,5 @@ class TestEncodeDecode(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(frame, EXPTECTED_FRAME)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()