refactor data models to use pydantic

.vscode/settings.json

@@ -11,9 +11,13 @@
         "test*.py"
     ],
     "python.testing.pytestEnabled": true,
-    "python.testing.unittestEnabled": true,
+    "python.testing.unittestEnabled": false,
     "python.formatting.provider": "black",
     "python.linting.mypyEnabled": true,
     "python.linting.mypyArgs": ["--config-file=${workspaceFolder}/mypy.ini"],
     "python.envFile": "${workspaceFolder}/.env",
+    "editor.formatOnSave": true,
+    "python.testing.pytestArgs": [
+        "src"
+    ]
 }

dev-requirements.txt

@@ -1,2 +1,4 @@
 pylint==2.17.2 
 pytest==7.3.1
+mypy==1.6.1
+types-requests==2.*

mypy.ini

@@ -0,0 +1,8 @@
+# Global options:
+
+[mypy]
+python_version = 3.10
+disallow_untyped_defs = True
+python_executable = venv/bin/python3.10
+follow_imports = silent
+plugins = pydantic.mypy

requirements.txt

@@ -1,2 +1,3 @@
-requests==2.28.2
-websocket-client==1.5.1
+requests==2.*
+websocket-client==1.5.1
+pydantic==2.4.2

src/dirigera/devices/base_ikea_model.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel, ConfigDict, alias_generators
+
+
+class BaseIkeaModel(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=alias_generators.to_camel, arbitrary_types_allowed=True
+    )

src/dirigera/devices/blinds.py

@@ -1,69 +1,44 @@
-from dataclasses import dataclass
+from __future__ import annotations
 from typing import Any, Optional, Dict
-
-from .device import Device
+from .device import Attributes, Device
 from ..hub.abstract_smart_home_hub import AbstractSmartHomeHub
 
 
-@dataclass
+class BlindAttributes(Attributes):
+    blinds_current_level: Optional[int] = None
+    blinds_target_level: Optional[int] = None
+    blinds_state: Optional[str] = None
+
+
 class Blind(Device):
     dirigera_client: AbstractSmartHomeHub
-    is_reachable: bool
-    current_level: Optional[int]
-    target_level: Optional[int]
-    state: Optional[str]
+    attributes: BlindAttributes
 
-    def refresh(self) -> None:
-        data = self.dirigera_client.get(route=f"/devices/{self.device_id}")
-        attributes: Dict[str, Any] = data["attributes"]
-        self.device_id = data["id"]
-        self.is_reachable = data["isReachable"]
-        self.custom_name = attributes["customName"]
-        self.current_level = attributes.get("blindsCurrentLevel")
-        self.target_level = attributes.get("blindsTargetLevel")
-        self.state = attributes.get("blindsState")
-        self.firmware_version = attributes.get("firmwareVersion")
-        self.room_id = data["room"]["id"]
-        self.room_name = data["room"]["name"]
-        self.can_receive = data["capabilities"]["canReceive"]
+    def reload(self) -> Blind:
+        data = self.dirigera_client.get(route=f"/devices/{self.id}")
+        return Blind(dirigeraClient=self.dirigera_client, **data)
 
     def set_name(self, name: str) -> None:
-        if "customName" not in self.can_receive:
+        if "customName" not in self.capabilities.can_receive:
             raise AssertionError("This blind does not support the customName function")
 
         data = [{"attributes": {"customName": name}}]
-        self.dirigera_client.patch(route=f"/devices/{self.device_id}", data=data)
-        self.custom_name = name
+        self.dirigera_client.patch(route=f"/devices/{self.id}", data=data)
+        self.attributes.custom_name = name
 
     def set_target_level(self, target_level: int) -> None:
-        if "blindsTargetLevel" not in self.can_receive:
-            raise AssertionError("This blind does not support the target level function")
+        if "blindsTargetLevel" not in self.capabilities.can_receive:
+            raise AssertionError(
+                "This blind does not support the target level function"
+            )
 
         if target_level < 0 or target_level > 100:
             raise AssertionError("target_level must be a value between 0 and 100")
 
         data = [{"attributes": {"blindsTargetLevel": target_level}}]
-        self.dirigera_client.patch(route=f"/devices/{self.device_id}", data=data)
-        self.target_level = target_level
+        self.dirigera_client.patch(route=f"/devices/{self.id}", data=data)
+        self.attributes.blinds_target_level = target_level
 
 
 def dict_to_blind(data: Dict[str, Any], dirigera_client: AbstractSmartHomeHub) -> Blind:
-    attributes: Dict[str, Any] = data["attributes"]
-
-    return Blind(
-        dirigera_client=dirigera_client,
-        device_id=data["id"],
-        is_reachable=data["isReachable"],
-        custom_name=attributes["customName"],
-        target_level=attributes["blindsTargetLevel"],
-        current_level=attributes["blindsCurrentLevel"],
-        state=attributes["blindsState"],
-        can_receive=data["capabilities"]["canReceive"],
-        room_id=data["room"]["id"],
-        room_name=data["room"]["name"],
-        firmware_version=attributes.get("firmwareVersion"),
-        hardware_version=attributes.get("hardwareVersion"),
-        model=attributes.get("model"),
-        manufacturer=attributes.get("manufacturer"),
-        serial_number=attributes.get("serialNumber"),
-    )
+    return Blind(dirigeraClient=dirigera_client, **data)

src/dirigera/devices/controller.py

@@ -1,61 +1,34 @@
-from dataclasses import dataclass
 from typing import Any, Optional, Dict
 
-from .device import Device
+from .device import Attributes, Device
 from ..hub.abstract_smart_home_hub import AbstractSmartHomeHub
 
 
-@dataclass
+class ControllerAttributes(Attributes):
+    is_on: bool
+    battery_percentage: Optional[int] = None
+
+
 class Controller(Device):
     dirigera_client: AbstractSmartHomeHub
-    is_reachable: bool
-    is_on: bool
-    battery_percentage: Optional[int]
-
-    def refresh(self):
-        data = self.dirigera_client.get(route=f"/devices/{self.device_id}")
-        attributes: Dict[str, Any] = data["attributes"]
-        self.device_id = data["id"]
-        self.is_reachable = data["isReachable"]
-        self.custom_name = attributes["customName"]
-        self.is_on = attributes["isOn"]
-        self.battery_percentage = attributes.get("batteryPercentage")
-        self.firmware_version = attributes.get("firmwareVersion")
-        self.room_id = data["room"]["id"]
-        self.can_receive = data["capabilities"]["canReceive"]
-        self.room_name = data["room"]["name"]
+    attributes: ControllerAttributes
+
+    def reload(self) -> None:
+        data = self.dirigera_client.get(route=f"/devices/{self.id}")
+        return Controller(dirigeraClient=self.dirigera_client, **data)
 
     def set_name(self, name: str) -> None:
-        if "customName" not in self.can_receive:
+        if "customName" not in self.capabilities.can_receive:
             raise AssertionError(
                 "This controller does not support the set_name function"
             )
 
         data = [{"attributes": {"customName": name}}]
-        self.dirigera_client.patch(
-            route=f"/devices/{self.device_id}", data=data
-        )
-        self.custom_name = name
+        self.dirigera_client.patch(route=f"/devices/{self.id}", data=data)
+        self.attributes.custom_name = name
 
 
 def dict_to_controller(
     data: Dict[str, Any], dirigera_client: AbstractSmartHomeHub
 ) -> Controller:
-    attributes: Dict[str, Any] = data["attributes"]
-
-    return Controller(
-        dirigera_client=dirigera_client,
-        device_id=data["id"],
-        is_reachable=data["isReachable"],
-        custom_name=attributes["customName"],
-        is_on=attributes["isOn"],
-        battery_percentage=attributes.get("batteryPercentage"),
-        can_receive=data["capabilities"]["canReceive"],
-        room_id=data["room"]["id"],
-        room_name=data["room"]["name"],
-        firmware_version=attributes.get("firmwareVersion"),
-        hardware_version=attributes.get("hardwareVersion"),
-        model=attributes.get("model"),
-        manufacturer=attributes.get("manufacturer"),
-        serial_number=attributes.get("serialNumber"),
-    )
+    return Controller(dirigeraClient=dirigera_client, **data)

src/dirigera/devices/device.py

@@ -1,6 +1,9 @@
-from dataclasses import dataclass
+from __future__ import annotations
+import datetime
 from enum import Enum
-from typing import Optional, List
+from typing import Any, Dict, Optional, List
+
+from pydantic import BaseModel, ConfigDict, alias_generators
 
 
 class StartupEnum(Enum):
@@ -10,15 +13,63 @@ class StartupEnum(Enum):
     START_TOGGLE = "startToggle"
 
 
-@dataclass
-class Device:
-    device_id: str
+class Attributes(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=alias_generators.to_camel, arbitrary_types_allowed=True
+    )
+
     custom_name: str
-    room_id: str
-    room_name: str
-    firmware_version: Optional[str]
-    hardware_version: Optional[str]
-    model: Optional[str]
-    manufacturer: Optional[str]
-    serial_number: Optional[str]
+    model: str
+    manufacturer: str
+    firmware_version: str
+    hardware_version: str
+    serial_number: Optional[str] = None
+    product_code: Optional[str] = None
+    ota_status: Optional[str] = None
+    ota_state: Optional[str] = None
+    ota_progress: Optional[int] = None
+    ota_policy: Optional[str] = None
+    ota_schedule_start: Optional[datetime.time] = None
+    ota_schedule_end: Optional[datetime.time] = None
+
+
+class Capabilities(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=alias_generators.to_camel, arbitrary_types_allowed=True
+    )
+
+    can_send: List[str]
     can_receive: List[str]
+
+
+class Room(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=alias_generators.to_camel, arbitrary_types_allowed=True
+    )
+
+    id: str
+    name: str
+    color: str
+    icon: str
+
+
+class Device(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=alias_generators.to_camel, arbitrary_types_allowed=True
+    )
+
+    id: str
+    type: str
+    device_type: str
+    created_at: datetime.datetime
+    is_reachable: bool
+    last_seen: datetime.datetime
+    attributes: Attributes
+    capabilities: Capabilities
+    room: Room
+    device_set: list[str]
+    remote_links: list[str]
+    is_hidden: Optional[bool] = None
+
+    def _reload(self, data: Dict[str, Any]) -> Device:
+        return Device(**data)

src/dirigera/devices/environment_sensor.py

@@ -1,62 +1,37 @@
-from dataclasses import dataclass
+from __future__ import annotations
 from typing import Any, Dict
-
-from .device import Device
+from .device import Attributes, Device
 from ..hub.abstract_smart_home_hub import AbstractSmartHomeHub
 
 
-@dataclass
+class EnvironmentSensorAttributes(Attributes):
+    current_temperature: int
+    current_r_h: int
+    current_p_m25: int
+    max_measured_p_m25: int
+    min_measured_p_m25: int
+    voc_index: int
+
+
 class EnvironmentSensor(Device):
     dirigera_client: AbstractSmartHomeHub
-    is_reachable: bool
-    current_temperature: str
-    current_rh: int
-    current_pm25: int
-    max_measured_pm25: int
-    min_measured_pm25: int
-    voc_index: int
+    attributes: EnvironmentSensorAttributes
 
-    def refresh(self) -> None:
-        data = self.dirigera_client.get(route=f"/devices/{self.device_id}")
-        attributes: Dict[str, Any] = data["attributes"]
-        self.firmware_version = attributes["firmwareVersion"]
-        self.current_temperature = attributes["currentTemperature"]
-        self.current_rh = attributes["currentRH"]
-        self.current_pm25 = attributes["currentPM25"]
-        self.voc_index = attributes["vocIndex"]
-        self.room_id = data["room"]["id"]
-        self.room_name = data["room"]["name"]
+    def reload(self) -> EnvironmentSensor:
+        data = self.dirigera_client.get(route=f"/devices/{self.id}")
+        return EnvironmentSensor(dirigeraClient=self.dirigera_client, **data)
 
     def set_name(self, name: str) -> None:
-        if "customName" not in self.can_receive:
+        if "customName" not in self.capabilities.can_receive:
             raise AssertionError("This sensor does not support the set_name function")
 
         data = [{"attributes": {"customName": name}}]
-        self.dirigera_client.patch(route=f"/devices/{self.device_id}", data=data)
-        self.custom_name = name
+        self.dirigera_client.patch(route=f"/devices/{self.id}", data=data)
+        self.attributes.custom_name = name
 
 
 def dict_to_environment_sensor(
     data: Dict[str, Any], dirigera_client: AbstractSmartHomeHub
-):
-    attributes: Dict[str, Any] = data["attributes"]
-    return EnvironmentSensor(
-        dirigera_client=dirigera_client,
-        device_id=data["id"],
-        is_reachable=data["isReachable"],
-        custom_name=attributes["customName"],
-        firmware_version=attributes.get("firmwareVersion"),
-        hardware_version=attributes.get("hardwareVersion"),
-        model=attributes.get("model"),
-        manufacturer=attributes.get("manufacturer"),
-        serial_number=attributes.get("serialNumber"),
-        current_temperature=attributes["currentTemperature"],
-        current_rh=attributes["currentRH"],
-        current_pm25=attributes["currentPM25"],
-        max_measured_pm25=attributes["maxMeasuredPM25"],
-        min_measured_pm25=attributes["minMeasuredPM25"],
-        voc_index=attributes["vocIndex"],
-        room_id=data["room"]["id"],
-        room_name=data["room"]["name"],
-        can_receive=data["capabilities"]["canReceive"],
-    )
+) -> EnvironmentSensor:
+    print(data)
+    return EnvironmentSensor(dirigeraClient=dirigera_client, **data)