Repository: mill1000/midea-msmart
Pull request #31 (2023-08-18 03:43:28)

Refactor class names to match PEP8 recommendations
This is a breaking change hopefully for obvious reasons


Patch:
diff --git a/example.py b/example.py
index 4b03cef..aa36e59 100644
--- a/example.py
+++ b/example.py
@@ -1,7 +1,7 @@
 import asyncio
 import logging
 
-from msmart.device import air_conditioning as ac
+from msmart.device import AirConditioner as AC
 from msmart.discover import Discover
 
 logging.basicConfig(level=logging.INFO)
@@ -26,7 +26,7 @@ async def main():
 
     # Manually construct the device
     #  - See midea-discover to read ID, token and key
-    device = ac(ip=DEVICE_IP, port=6444, device_id=int(DEVICE_ID))
+    device = AC(ip=DEVICE_IP, port=6444, device_id=int(DEVICE_ID))
     if DEVICE_TOKEN and DEVICE_KEY:
         await device.authenticate(DEVICE_TOKEN, DEVICE_KEY)
 
@@ -60,7 +60,7 @@ async def main():
     device.power_state = True
     device.prompt_tone = False
     device.target_temperature = 25
-    device.operational_mode = ac.operational_mode_enum.cool
+    device.operational_mode = AC.OperationalMode.COOL
     await device.apply()
 
     print({
diff --git a/msmart/base_command.py b/msmart/base_command.py
index 1781440..0ba30f6 100644
--- a/msmart/base_command.py
+++ b/msmart/base_command.py
@@ -2,15 +2,15 @@ import logging
 from abc import ABC, abstractmethod
 
 import msmart.crc8 as crc8
-from msmart.const import FRAME_TYPE
+from msmart.const import FrameType
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class command(ABC):
+class Command(ABC):
     _message_id = 0
 
-    def __init__(self, device_type=0xAC, frame_type=FRAME_TYPE.Request):
+    def __init__(self, device_type=0xAC, frame_type=FrameType.REQUEST):
         self.device_type = device_type
         self.frame_type = frame_type
         self.protocol_version = 0
@@ -51,7 +51,7 @@ class command(ABC):
         frame = header + payload_crc
 
         # Calculate total frame checksum
-        frame.append(command.checksum(frame[1:]))
+        frame.append(Command.checksum(frame[1:]))
 
         _LOGGER.debug("Frame data: %s", frame.hex())
 
@@ -63,20 +63,10 @@ class command(ABC):
 
     @property
     def message_id(self):
-        command._message_id += 1
-        return command._message_id & 0xFF
+        Command._message_id += 1
+        return Command._message_id & 0xFF
 
     @property
     @abstractmethod
     def payload(self):
         return bytes()
-
-
-class set_customize_command(command):
-    def __init__(self, device_type, frame_type, customize_cmd,):
-        super().__init__(device_type, frame_type=frame_type.Request)
-        self.customize_cmd = customize_cmd
-
-    @property
-    def payload(self):
-        return bytearray.fromhex(self.customize_cmd)
diff --git a/msmart/device/base.py b/msmart/base_device.py
similarity index 98%
rename from msmart/device/base.py
rename to msmart/base_device.py
index 2de7511..ec19719 100644
--- a/msmart/device/base.py
+++ b/msmart/base_device.py
@@ -8,7 +8,7 @@ from msmart.types import Key, Token
 _LOGGER = logging.getLogger(__name__)
 
 
-class device():
+class Device():
 
     def __init__(self, *, ip: str, port: int, device_id: int, **kwargs) -> None:
         self._ip = ip
@@ -112,7 +112,7 @@ class device():
 
     def __str__(self) -> str:
         return str({
-            "ip" : self.ip,
+            "ip": self.ip,
             "port": self.port,
             "id": self.id,
             "online": self.online,
diff --git a/msmart/const.py b/msmart/const.py
index c9160b4..8761423 100644
--- a/msmart/const.py
+++ b/msmart/const.py
@@ -31,9 +31,9 @@ class DeviceType(IntEnum):
     AIR_CONDITIONER = 0xAC
 
 
-class FRAME_TYPE(IntEnum):
-    Unknown = 0
-    Set = 0x02
-    Request = 0x03
-    Report = 0x04
-    ABNL_REPORT = 0x06
+class FrameType(IntEnum):
+    UNKNOWN = 0
+    SET = 0x02
+    REQUEST = 0x03
+    REPORT = 0x04
+    ABNORMAL_REPORT = 0x06
diff --git a/msmart/device/AC/command.py b/msmart/device/AC/command.py
index 1017106..c3998f7 100644
--- a/msmart/device/AC/command.py
+++ b/msmart/device/AC/command.py
@@ -5,8 +5,8 @@ from collections import namedtuple
 from enum import IntEnum
 
 import msmart.crc8 as crc8
-from msmart.base_command import command
-from msmart.const import FRAME_TYPE
+from msmart.base_command import Command
+from msmart.const import FrameType
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -16,44 +16,44 @@ class InvalidResponseException(Exception):
 
 
 class ResponseId(IntEnum):
-    State = 0xC0
-    Capabilities = 0xB5
+    STATE = 0xC0
+    CAPABILITIES = 0xB5
 
 
 class CapabilityId(IntEnum):
-    IndoorHumidity = 0x0015
-    SilkyCool = 0x0018
-    SmartEye = 0x0030
-    WindOnMe = 0x0032
-    WindOffMe = 0x0033
-    ActiveClean = 0x0039
-    OneKeyNoWindOnMe = 0x0042
-    BreezeControl = 0x0043
-    FanSpeedControl = 0x0210
-    PresetEco = 0x0212
-    PresetFreezeProtection = 0x0213
-    Modes = 0x0214
-    SwingModes = 0x0215
-    Power = 0x0216
-    Nest = 0x0217
-    AuxElectricHeat = 0x0219
-    PresetTurbo = 0x021A
-    Humidity = 0x021F
-    UnitChangeable = 0x0222
-    LightControl = 0x0224
-    Temperatures = 0x0225
-    Buzzer = 0x022C
-
-
-class temperature_type(IntEnum):
-    Unknown = 0
-    Indoor = 0x2
-    Outdoor = 0x3
-
-
-class get_capabilities_command(command):
+    INDOOR_HUMIDITY = 0x0015
+    SILKY_COOL = 0x0018
+    SMART_EYE = 0x0030
+    WIND_ON_ME = 0x0032
+    WIND_OFF_ME = 0x0033
+    ACTIVE_CLEAN = 0x0039
+    ONE_KEY_NO_WIND_ON_ME = 0x0042
+    BREEZE_CONTROL = 0x0043
+    FAN_SPEED_CONTROL = 0x0210
+    PRESET_ECO = 0x0212
+    PRESET_FREEZE_PROTECTION = 0x0213
+    MODES = 0x0214
+    SWING_MODES = 0x0215
+    POWER = 0x0216
+    NEST = 0x0217
+    AUX_ELECTRIC_HEAT = 0x0219
+    PRESET_TURBO = 0x021A
+    HUMIDITY = 0x021F
+    UNIT_CHANGEABLE = 0x0222
+    LIGHT_CONTROL = 0x0224
+    TEMPERATURES = 0x0225
+    BUZZER = 0x022C
+
+
+class TemperatureType(IntEnum):
+    UNKNOWN = 0
+    INDOOR = 0x2
+    OUTDOOR = 0x3
+
+
+class GetCapabilitiesCommand(Command):
     def __init__(self, device_type):
-        super().__init__(device_type, frame_type=FRAME_TYPE.Request)
+        super().__init__(device_type, frame_type=FrameType.REQUEST)
 
     @property
     def payload(self):
@@ -65,11 +65,11 @@ class get_capabilities_command(command):
         ])
 
 
-class get_state_command(command):
+class GetStateCommand(Command):
     def __init__(self, device_type):
-        super().__init__(device_type, frame_type=FRAME_TYPE.Request)
+        super().__init__(device_type, frame_type=FrameType.REQUEST)
 
-        self.temperature_type = temperature_type.Indoor
+        self.temperature_type = TemperatureType.INDOOR
 
     @property
     def payload(self):
@@ -89,9 +89,9 @@ class get_state_command(command):
         ])
 
 
-class set_state_command(command):
+class SetStateCommand(Command):
     def __init__(self, device_type):
-        super().__init__(device_type, frame_type=FRAME_TYPE.Set)
+        super().__init__(device_type, frame_type=FrameType.SET)
 
         self.beep_on = True
         self.power_on = False
@@ -163,10 +163,10 @@ class set_state_command(command):
         ])
 
 
-class toggle_display_command(command):
+class ToggleDisplayCommand(Command):
     def __init__(self, device_type):
         # For whatever reason, toggle display uses a request type...
-        super().__init__(device_type, frame_type=FRAME_TYPE.Request)
+        super().__init__(device_type, frame_type=FrameType.REQUEST)
 
     @property
     def payload(self):
@@ -183,7 +183,7 @@ class toggle_display_command(command):
         ])
 
 
-class response():
+class Response():
     def __init__(self, payload: memoryview):
         # Set ID and copy the payload
         self._id = payload[0]
@@ -200,7 +200,7 @@ class response():
     @staticmethod
     def validate(frame: memoryview):
         # Validate frame checksum
-        frame_checksum = command.checksum(frame[1:-1])
+        frame_checksum = Command.checksum(frame[1:-1])
         if frame_checksum != frame[-1]:
             raise InvalidResponseException(
                 f"Frame '{frame.hex()}' failed checksum. Received: 0x{frame[-1]:X}, Expected: 0x{frame_checksum:X}.")
@@ -210,7 +210,7 @@ class response():
 
         # Some devices use a CRC others seem to use a 2nd checksum
         payload_crc = crc8.calculate(payload[0:-1])
-        payload_checksum = command.checksum(payload[0:-1])
+        payload_checksum = Command.checksum(payload[0:-1])
 
         if payload_crc != payload[-1] and payload_checksum != payload[-1]:
             raise InvalidResponseException(
@@ -221,20 +221,20 @@ class response():
         # Build a memoryview of the frame for zero-copy slicing
         with memoryview(frame) as frame_mv:
             # Ensure frame is valid before parsing
-            response.validate(frame_mv)
+            Response.validate(frame_mv)
 
             # Parse frame depending on id
             response_id = frame_mv[10]
             payload = frame_mv[10:-2]
-            if response_id == ResponseId.State:
-                return state_response(payload)
-            elif response_id == ResponseId.Capabilities:
-                return capabilities_response(payload)
+            if response_id == ResponseId.STATE:
+                return StateResponse(payload)
+            elif response_id == ResponseId.CAPABILITIES:
+                return CapabilitiesResponse(payload)
             else:
-                return response(payload)
+                return Response(payload)
 
 
-class capabilities_response(response):
+class CapabilitiesResponse(Response):
     def __init__(self, payload: memoryview):
         super().__init__(payload)
 
@@ -260,53 +260,53 @@ class capabilities_response(response):
 
         # Create a map of capability ID to decoders
         capability_readers = {
-            CapabilityId.IndoorHumidity: reader("indoor_humidity", get_bool),
-            CapabilityId.SilkyCool: reader("silky_cool", get_value(1)),
-            CapabilityId.SmartEye:  reader("smart_eye", get_value(1)),
-            CapabilityId.WindOnMe:  reader("wind_on_me", get_value(1)),
-            CapabilityId.WindOffMe:  reader("wind_off_me", get_value(1)),
-            CapabilityId.ActiveClean:  reader("active_clean", get_value(1)),
-            CapabilityId.OneKeyNoWindOnMe: reader("one_key_no_wind_on_me", get_value(1)),
-            CapabilityId.BreezeControl: reader("breeze_control", get_value(1)),
+            CapabilityId.INDOOR_HUMIDITY: reader("indoor_humidity", get_bool),
+            CapabilityId.SILKY_COOL: reader("silky_cool", get_value(1)),
+            CapabilityId.SMART_EYE:  reader("smart_eye", get_value(1)),
+            CapabilityId.WIND_ON_ME:  reader("wind_on_me", get_value(1)),
+            CapabilityId.WIND_OFF_ME:  reader("wind_off_me", get_value(1)),
+            CapabilityId.ACTIVE_CLEAN:  reader("active_clean", get_value(1)),
+            CapabilityId.ONE_KEY_NO_WIND_ON_ME: reader("one_key_no_wind_on_me", get_value(1)),
+            CapabilityId.BREEZE_CONTROL: reader("breeze_control", get_value(1)),
             # Fan speed control always seems to return false, even if unit can
-            CapabilityId.FanSpeedControl: reader("fan_speed_control", get_no_value(1)),
-            CapabilityId.PresetEco: [
+            CapabilityId.FAN_SPEED_CONTROL: reader("fan_speed_control", get_no_value(1)),
+            CapabilityId.PRESET_ECO: [
                 reader("eco_mode", get_value(1)),
                 reader("eco_mode_2", get_value(2)),
             ],
-            CapabilityId.PresetFreezeProtection: reader("freeze_protection", get_value(1)),
-            CapabilityId.Modes: [
+            CapabilityId.PRESET_FREEZE_PROTECTION: reader("freeze_protection", get_value(1)),
+            CapabilityId.MODES: [
                 reader("heat_mode", lambda v: v == 1 or v == 2),
                 reader("cool_mode", lambda v: v != 2),
                 reader("dry_mode", lambda v: v < 2),
                 reader("auto_mode", lambda v: v < 3),
             ],
-            CapabilityId.SwingModes: [
+            CapabilityId.SWING_MODES: [
                 reader("swing_horizontal", lambda v: v == 1 or v == 3),
                 reader("swing_vertical", lambda v: v < 2),
             ],
-            CapabilityId.Power: [
+            CapabilityId.POWER: [
                 reader("power_cal", lambda v: v == 2 or v == 3),
                 reader("power_cal_setting", lambda v: v == 3),
             ],
-            CapabilityId.Nest: [
+            CapabilityId.NEST: [
                 reader("nest_check", lambda v: v == 1 or v == 2 or v == 4),
                 reader("nest_need_change", lambda v: v == 3 or v == 4),
             ],
-            CapabilityId.AuxElectricHeat: reader("aux_electric_heat", get_bool),
-            CapabilityId.PresetTurbo:  [
+            CapabilityId.AUX_ELECTRIC_HEAT: reader("aux_electric_heat", get_bool),
+            CapabilityId.PRESET_TURBO:  [
                 reader("turbo_heat", lambda v: v == 1 or v == 3),
                 reader("turbo_cool", lambda v: v < 2),
             ],
-            CapabilityId.Humidity:
+            CapabilityId.HUMIDITY:
             [
                 reader("humidity_auto_set", lambda v: v == 1 or v == 2),
                 reader("humidity_manual_set", lambda v: v == 2 or v == 3),
             ],
-            CapabilityId.UnitChangeable: reader("unit_changeable", get_value(0)),
-            CapabilityId.LightControl: reader("light_control", get_bool),
+            CapabilityId.UNIT_CHANGEABLE: reader("unit_changeable", get_value(0)),
+            CapabilityId.LIGHT_CONTROL: reader("light_control", get_bool),
             # Temperatures capability too complex to be handled here
-            CapabilityId.Buzzer:  reader("buzzer", get_bool),
+            CapabilityId.BUZZER:  reader("buzzer", get_bool),
         }
 
         count = payload[1]
@@ -353,7 +353,7 @@ class capabilities_response(response):
                     # Apply the single reader
                     self._capabilities.update(apply(reader, value))
 
-            elif capability_id == CapabilityId.Temperatures:
+            elif capability_id == CapabilityId.TEMPERATURES:
                 # Skip if capability size is too small
                 if size < 6:
                     continue
@@ -429,7 +429,7 @@ class capabilities_response(response):
         return self._capabilities.get("freeze_protection", False)
 
 
-class state_response(response):
+class StateResponse(Response):
     def __init__(self, payload: memoryview):
         super().__init__(payload)
 
diff --git a/msmart/device/AC/appliance.py b/msmart/device/AC/device.py
similarity index 76%
rename from msmart/device/AC/appliance.py
rename to msmart/device/AC/device.py
index 61c2867..c737132 100644
--- a/msmart/device/AC/appliance.py
+++ b/msmart/device/AC/device.py
@@ -2,14 +2,14 @@ import logging
 from enum import IntEnum
 from typing import Any, List, Optional
 
+from msmart.base_device import Device
 from msmart.const import DeviceType
-from msmart.device.base import device
 
-from .command import (InvalidResponseException, ResponseId,
-                      capabilities_response, get_capabilities_command,
-                      get_state_command)
-from .command import response as base_response
-from .command import set_state_command, state_response, toggle_display_command
+from .command import (CapabilitiesResponse, GetCapabilitiesCommand,
+                      GetStateCommand, InvalidResponseException)
+from .command import Response as base_response
+from .command import (ResponseId, SetStateCommand, StateResponse,
+                      ToggleDisplayCommand)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -29,39 +29,39 @@ class IntEnumHelper(IntEnum):
             return default
 
 
-class air_conditioning(device):
+class AirConditioner(Device):
 
-    class fan_speed_enum(IntEnumHelper):
-        Auto = 102
-        Full = 100
-        High = 80
-        Medium = 60
-        Low = 40
-        Silent = 20
+    class FanSpeed(IntEnumHelper):
+        AUTO = 102
+        FULL = 100
+        HIGH = 80
+        MEDIUM = 60
+        LOW = 40
+        SILENT = 20
 
         @classmethod
-        def get(cls, value, default=Auto) -> IntEnum:
+        def get(cls, value, default=AUTO) -> IntEnum:
             return super().get(value, default)
 
-    class operational_mode_enum(IntEnumHelper):
-        auto = 1
-        cool = 2
-        dry = 3
-        heat = 4
-        fan_only = 5
+    class OperationalMode(IntEnumHelper):
+        AUTO = 1
+        COOL = 2
+        DRY = 3
+        HEAT = 4
+        FAN_ONLY = 5
 
         @classmethod
-        def get(cls, value, default=fan_only) -> IntEnum:
+        def get(cls, value, default=FAN_ONLY) -> IntEnum:
             return super().get(value, default)
 
-    class swing_mode_enum(IntEnumHelper):
-        Off = 0x0
-        Vertical = 0xC
-        Horizontal = 0x3
-        Both = 0xF
+    class SwingMode(IntEnumHelper):
+        OFF = 0x0
+        VERTICAL = 0xC
+        HORIZONTAL = 0x3
+        BOTH = 0xF
 
         @classmethod
-        def get(cls, value, default=Off) -> IntEnum:
+        def get(cls, value, default=OFF) -> IntEnum:
             return super().get(value, default)
 
     def __init__(self, ip: str, device_id: int,  port: int, **kwargs) -> None:
@@ -77,9 +77,9 @@ class air_conditioning(device):
         self._prompt_tone = False
         self._power_state = False
         self._target_temperature = 17.0
-        self._operational_mode = air_conditioning.operational_mode_enum.auto
-        self._fan_speed = air_conditioning.fan_speed_enum.Auto
-        self._swing_mode = air_conditioning.swing_mode_enum.Off
+        self._operational_mode = AirConditioner.OperationalMode.AUTO
+        self._fan_speed = AirConditioner.FanSpeed.AUTO
+        self._swing_mode = AirConditioner.SwingMode.OFF
         self._eco_mode = False
         self._turbo_mode = False
         self._freeze_protection_mode = False
@@ -89,8 +89,8 @@ class air_conditioning(device):
         self._filter_alert = False
 
         # Support all known modes initially
-        self._supported_op_modes = air_conditioning.operational_mode_enum.list()
-        self._supported_swing_modes = air_conditioning.swing_mode_enum.list()
+        self._supported_op_modes = AirConditioner.OperationalMode.list()
+        self._supported_swing_modes = AirConditioner.SwingMode.list()
         self._supports_eco = True
         self._supports_turbo = True
         self._supports_freeze_protection_mode = True
@@ -104,21 +104,21 @@ class air_conditioning(device):
         self._outdoor_temperature = None
 
     async def get_capabilities(self) -> None:
-        cmd = get_capabilities_command(self.type)
+        cmd = GetCapabilitiesCommand(self.type)
         await self.send_command(cmd)
 
     async def toggle_display(self) -> None:
         if not self._supports_display_control:
             _LOGGER.warning("Device is not capable of display control.")
 
-        cmd = toggle_display_command(self.type)
+        cmd = ToggleDisplayCommand(self.type)
         await self.send_command(cmd, True)
 
         # Force a refresh to get the updated display state
         await self.refresh()
 
     async def refresh(self):
-        cmd = get_state_command(self.type)
+        cmd = GetStateCommand(self.type)
         await self.send_command(cmd)
 
     async def send_command(self, command, ignore_response=False) -> None:
@@ -144,9 +144,9 @@ class air_conditioning(device):
 
             self._supported = True
 
-            if response.id == ResponseId.State:
+            if response.id == ResponseId.STATE:
                 self.update(response)
-            elif response.id == ResponseId.Capabilities:
+            elif response.id == ResponseId.CAPABILITIES:
                 self.update_capabilities(response)
             else:
                 _LOGGER.debug("Ignored unknown response from %s:%d: %s",
@@ -175,7 +175,7 @@ class air_conditioning(device):
             if self._freeze_protection_mode and not self._supports_freeze_protection_mode:
                 _LOGGER.warning("Device is not capable of freeze protection.")
 
-            cmd = set_state_command(self.type)
+            cmd = SetStateCommand(self.type)
             cmd.beep_on = self._prompt_tone
             cmd.power_on = self._power_state
             cmd.target_temperature = self._target_temperature
@@ -192,17 +192,17 @@ class air_conditioning(device):
             self._updating = False
             self._defer_update = False
 
-    def update(self, res: state_response) -> None:
+    def update(self, res: StateResponse) -> None:
         self._power_state = res.power_on
 
         self._target_temperature = res.target_temperature
-        self._operational_mode = air_conditioning.operational_mode_enum.get(
+        self._operational_mode = AirConditioner.OperationalMode.get(
             res.operational_mode)
 
-        self._fan_speed = air_conditioning.fan_speed_enum.get(
+        self._fan_speed = AirConditioner.FanSpeed.get(
             res.fan_speed)
 
-        self._swing_mode = air_conditioning.swing_mode_enum.get(
+        self._swing_mode = AirConditioner.SwingMode.get(
             res.swing_mode)
 
         self._eco_mode = res.eco_mode
@@ -221,28 +221,28 @@ class air_conditioning(device):
         # self._on_timer = res.on_timer
         # self._off_timer = res.off_timer
 
-    def update_capabilities(self, res: capabilities_response) -> None:
+    def update_capabilities(self, res: CapabilitiesResponse) -> None:
         # Build list of supported operation modes
-        op_modes = [air_conditioning.operational_mode_enum.fan_only]
+        op_modes = [AirConditioner.OperationalMode.FAN_ONLY]
         if res.dry_mode:
-            op_modes.append(air_conditioning.operational_mode_enum.dry)
+            op_modes.append(AirConditioner.OperationalMode.DRY)
         if res.cool_mode:
-            op_modes.append(air_conditioning.operational_mode_enum.cool)
+            op_modes.append(AirConditioner.OperationalMode.COOL)
         if res.heat_mode:
-            op_modes.append(air_conditioning.operational_mode_enum.heat)
+            op_modes.append(AirConditioner.OperationalMode.HEAT)
         if res.auto_mode:
-            op_modes.append(air_conditioning.operational_mode_enum.auto)
+            op_modes.append(AirConditioner.OperationalMode.AUTO)
 
         self._supported_op_modes = op_modes
 
         # Build list of supported swing modes
-        swing_modes = [air_conditioning.swing_mode_enum.Off]
+        swing_modes = [AirConditioner.SwingMode.OFF]
         if res.swing_horizontal:
-            swing_modes.append(air_conditioning.swing_mode_enum.Horizontal)
+            swing_modes.append(AirConditioner.SwingMode.HORIZONTAL)
         if res.swing_vertical:
-            swing_modes.append(air_conditioning.swing_mode_enum.Vertical)
+            swing_modes.append(AirConditioner.SwingMode.VERTICAL)
         if res.swing_both:
-            swing_modes.append(air_conditioning.swing_mode_enum.Both)
+            swing_modes.append(AirConditioner.SwingMode.BOTH)
 
         self._supported_swing_modes = swing_modes
 
@@ -286,31 +286,31 @@ class air_conditioning(device):
         self._target_temperature = temperature_celsius
 
     @property
-    def operational_mode(self) -> operational_mode_enum:
+    def operational_mode(self) -> OperationalMode:
         return self._operational_mode
 
     @operational_mode.setter
-    def operational_mode(self, mode: operational_mode_enum) -> None:
+    def operational_mode(self, mode: OperationalMode) -> None:
         if self._updating:
             self._defer_update = True
         self._operational_mode = mode
 
     @property
-    def fan_speed(self) -> fan_speed_enum:
+    def fan_speed(self) -> FanSpeed:
         return self._fan_speed
 
     @fan_speed.setter
-    def fan_speed(self, speed: fan_speed_enum) -> None:
+    def fan_speed(self, speed: FanSpeed) -> None:
         if self._updating:
             self._defer_update = True
         self._fan_speed = speed
 
     @property
-    def swing_mode(self) -> swing_mode_enum:
+    def swing_mode(self) -> SwingMode:
         return self._swing_mode
 
     @swing_mode.setter
-    def swing_mode(self, mode: swing_mode_enum) -> None:
+    def swing_mode(self, mode: SwingMode) -> None:
         if self._updating:
             self._defer_update = True
         self._swing_mode = mode
@@ -394,11 +394,11 @@ class air_conditioning(device):
         return self._off_timer
 
     @property
-    def supported_operation_modes(self) -> List[operational_mode_enum]:
+    def supported_operation_modes(self) -> List[OperationalMode]:
         return self._supported_op_modes
 
     @property
-    def supported_swing_modes(self) -> List[swing_mode_enum]:
+    def supported_swing_modes(self) -> List[SwingMode]:
         return self._supported_swing_modes
 
     @property
@@ -406,5 +406,5 @@ class air_conditioning(device):
         return self._min_target_temperature
 
     @property
-    def max_target_temperature(self)-> Optional[int] :
+    def max_target_temperature(self) -> Optional[int]:
         return self._max_target_temperature
diff --git a/msmart/device/__init__.py b/msmart/device/__init__.py
index ebb07e3..b449937 100644
--- a/msmart/device/__init__.py
+++ b/msmart/device/__init__.py
@@ -1,2 +1,3 @@
-from .AC.appliance import air_conditioning
-from .base import device
+from msmart.base_device import Device
+
+from .AC.device import AirConditioner
diff --git a/msmart/discover.py b/msmart/discover.py
index fa264df..0d65b7b 100644
--- a/msmart/discover.py
+++ b/msmart/discover.py
@@ -10,7 +10,7 @@ from msmart.cloud import Cloud
 from msmart.const import (DEVICE_INFO_MSG, DISCOVERY_MSG,
                           OPEN_MIDEA_APP_ACCOUNT, OPEN_MIDEA_APP_PASSWORD,
                           DeviceType)
-from msmart.device import air_conditioning, device
+from msmart.device import AirConditioner, Device
 from msmart.lan import Security
 
 _LOGGER = logging.getLogger(__name__)
@@ -141,7 +141,7 @@ class Discover:
         interface=None,
         account=None,
         password=None,
-    ) -> [device]:
+    ) -> [Device]:
         """Discover devices via broadcast."""
 
         Discover._set_cloud_credentials(account, password)
@@ -179,7 +179,7 @@ class Discover:
         cls,
         host,
         **kwargs
-    ) -> device:
+    ) -> Device:
         """Discover a single device by hostname or IP."""
 
         devices = await Discover.discover(target=host, **kwargs)
@@ -316,13 +316,13 @@ class Discover:
         """Get the device class from the device type."""
 
         if device_type == DeviceType.AIR_CONDITIONER:
-            return air_conditioning
+            return AirConditioner
 
         # Unknown type return generic device
-        return device
+        return Device
 
     @classmethod
-    async def _authenticate_device(cls, dev: device):
+    async def _authenticate_device(cls, dev: Device):
         """Attempt to authenticate a V3 device."""
 
         # Get cloud connection


Test patch:
diff --git a/msmart/device/AC/test_command.py b/msmart/device/AC/test_command.py
index efdbe06..c1ebb31 100644
--- a/msmart/device/AC/test_command.py
+++ b/msmart/device/AC/test_command.py
@@ -1,6 +1,6 @@
 import unittest
 
-from .command import response
+from .command import Response
 
 
 class TestStateResponse(unittest.TestCase):
@@ -10,7 +10,7 @@ class TestStateResponse(unittest.TestCase):
                         msg=f"Object {obj} lacks attribute '{attr}'.")
 
     def _test_build_response(self, msg):
-        resp = response.construct(msg)
+        resp = Response.construct(msg)
         self.assertIsNotNone(resp)
 
         return resp
diff --git a/msmart/tests/test_discover.py b/msmart/tests/test_discover.py
index d7c7a36..6db5d10 100644
--- a/msmart/tests/test_discover.py
+++ b/msmart/tests/test_discover.py
@@ -1,7 +1,7 @@
 import unittest
 
 from msmart.const import DeviceType
-from msmart.device import air_conditioning as AC
+from msmart.device import AirConditioner as AC
 from msmart.discover import Discover