Implement pydantic for sockets

mypy.ini

@@ -146,3 +146,47 @@ disallow_untyped_defs = true
 no_implicit_optional = true
 warn_return_any = true
 warn_unreachable = true
+
+[mypy-pytradfri.device.socket]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unreachable = true
+
+[mypy-pytradfri.device.socket_control]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unreachable = true
+
+[mypy-pytradfri.device.air_purifier]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unreachable = true
+
+[mypy-pytradfri.device.air_purifier_control]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unreachable = true

pytradfri/device/__init__.py

@@ -29,6 +29,7 @@
 from .blind_control import BlindControl, BlindResponse
 from .light_control import LightControl
 from .signal_repeater_control import SignalRepeaterControl
+from .socket import SocketResponse
 from .socket_control import SocketControl
 
 
@@ -60,7 +61,7 @@ class DeviceResponse(ApiResourceResponse):
     signal_repeater_control: Optional[List[Dict[str, Any]]] = Field(
         alias=ROOT_SIGNAL_REPEATER
     )
-    socket_control: Optional[List[Dict[str, Any]]] = Field(alias=ATTR_SWITCH_PLUG)
+    socket_control: Optional[List[SocketResponse]] = Field(alias=ATTR_SWITCH_PLUG)
 
 
 class Device(ApiResource):

pytradfri/device/air_purifier_control.py

@@ -29,36 +29,36 @@ def air_purifiers(self) -> list[AirPurifier]:
         """Return air purifier objects of the air purifier control."""
         return [AirPurifier(self._device, i) for i in range(len(self.raw))]
 
-    def turn_off(self, *, index=0) -> Command:
+    def turn_off(self, *, index: int = 0) -> Command:
         """Turn the device off."""
         return self.set_value({ATTR_AIR_PURIFIER_MODE: 0}, index=index)
 
-    def turn_on_auto_mode(self, *, index=0) -> Command:
+    def turn_on_auto_mode(self, *, index: int = 0) -> Command:
         """Turn on auto mode."""
         return self.set_value(
             {ATTR_AIR_PURIFIER_MODE: ATTR_AIR_PURIFIER_MODE_AUTO}, index=index
         )
 
-    def set_fan_speed(self, mode: int, *, index=0) -> Command:
+    def set_fan_speed(self, mode: int, *, index: int = 0) -> Command:
         """Set the fan speed of the purifier."""
         self._value_validate(mode, RANGE_AIR_PURIFIER, "Air Purifier mode")
         return self.set_value({ATTR_AIR_PURIFIER_MODE: mode}, index=index)
 
-    def set_controls_locked(self, locked: bool, *, index=0) -> Command:
+    def set_controls_locked(self, locked: bool, *, index: int = 0) -> Command:
         """Set physical controls locked of the air purifier."""
 
         return self.set_value(
             {ATTR_AIR_PURIFIER_CONTROLS_LOCKED: 1 if locked else 0}, index=index
         )
 
-    def set_leds_off(self, leds_off: bool, *, index=0) -> Command:
+    def set_leds_off(self, leds_off: bool, *, index: int = 0) -> Command:
         """Set led's off/on of the air purifier."""
 
         return self.set_value(
             {ATTR_AIR_PURIFIER_LEDS_OFF: 1 if leds_off else 0}, index=index
         )
 
-    def set_value(self, value: dict[str, bool | int], *, index=0) -> Command:
+    def set_value(self, value: dict[str, bool | int], *, index: int = 0) -> Command:
         """Set values on air purifier control.
 
         Returns a Command.

pytradfri/device/socket.py

@@ -1,32 +1,46 @@
 """Represent a socket."""
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING
 
-from ..const import ATTR_DEVICE_STATE
+from pydantic import BaseModel, Field
+
+from ..const import ATTR_DEVICE_STATE, ATTR_ID
+
+
+class SocketResponse(BaseModel):
+    """Represent API response for a blind."""
+
+    id: int = Field(alias=ATTR_ID)
+    state: int = Field(alias=ATTR_DEVICE_STATE)
+
+
+if TYPE_CHECKING:
+    # avoid cyclic import at runtime.
+    from . import Device
 
 
 class Socket:
     """Represent a socket."""
 
-    def __init__(self, device, index):
+    def __init__(self, device: Device, index: int) -> None:
         """Create object of class."""
         self.device = device
         self.index = index
 
     @property
-    def state(self):
+    def state(self) -> bool:
         """State."""
-        return self.raw.get(ATTR_DEVICE_STATE) == 1
+        return self.raw.state == 1
 
     @property
-    def raw(self) -> dict[str, Any]:
+    def raw(self) -> SocketResponse:
         """Return raw data that it represents."""
         socket_control_response = self.device.raw.socket_control
         assert socket_control_response is not None
         return socket_control_response[self.index]
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return representation of class object."""
         state = "on" if self.state else "off"
         return f"<Socket #{self.index} - name: {self.device.name}, state: {state}>"

pytradfri/device/socket_control.py

@@ -1,34 +1,32 @@
 """Class to control the sockets."""
 from __future__ import annotations
 
-from typing import Any
-
 from ..command import Command
 from ..const import ATTR_DEVICE_STATE, ATTR_SWITCH_PLUG
 from .base_controller import BaseController
-from .socket import Socket
+from .socket import Socket, SocketResponse
 
 
 class SocketControl(BaseController):
     """Class to control the sockets."""
 
     @property
-    def raw(self) -> list[dict[str, Any]]:
+    def raw(self) -> list[SocketResponse]:
         """Return raw data that it represents."""
         socket_control_response = self._device.raw.socket_control
         assert socket_control_response is not None
         return socket_control_response
 
     @property
-    def sockets(self):
+    def sockets(self) -> list[Socket]:
         """Return socket objects of the socket control."""
         return [Socket(self._device, i) for i in range(len(self.raw))]
 
-    def set_state(self, state, *, index=0):
+    def set_state(self, state: int, *, index: int = 0) -> Command:
         """Set state of a socket."""
-        return self.set_values({ATTR_DEVICE_STATE: int(state)}, index=index)
+        return self.set_values({ATTR_DEVICE_STATE: state}, index=index)
 
-    def set_values(self, values, *, index=0):
+    def set_values(self, values: dict[str, int], *, index: int = 0) -> Command:
         """Set values on socket control.
 
         Returns a Command.

tests/test_device.py

@@ -336,15 +336,19 @@ def test_socket_value_setting(function_name, comment, test_input, expected_resul
 
 def test_socket_state_off():
     """Test socket off."""
-    socket = Device(deepcopy(OUTLET)).socket_control.sockets[0]
-    socket.raw[ATTR_DEVICE_STATE] = 0
+    socket_response = deepcopy(OUTLET)
+    socket_response[ATTR_SWITCH_PLUG][0][ATTR_DEVICE_STATE] = 0
+
+    socket = Device(socket_response).socket_control.sockets[0]
     assert socket.state is False
 
 
 def test_socket_state_on():
     """Test socket on."""
-    socket = Device(deepcopy(OUTLET)).socket_control.sockets[0]
-    socket.raw[ATTR_DEVICE_STATE] = 1
+    socket_response = deepcopy(OUTLET)
+    socket_response[ATTR_SWITCH_PLUG][0][ATTR_DEVICE_STATE] = 1
+
+    socket = Device(socket_response).socket_control.sockets[0]
     assert socket.state is True