From bbd7b7d7f18aec27b78e21abe283182b3e7d39de Mon Sep 17 00:00:00 2001
From: Simone Chemelli <simone.chemelli@gmail.com>
Date: Tue, 4 Jun 2024 14:22:56 +0000
Subject: [PATCH] chore: typing dc

---
 midealocal/devices/dc/__init__.py | 10 ++++++----
 midealocal/devices/dc/message.py  | 27 ++++++++++++++-------------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/midealocal/devices/dc/__init__.py b/midealocal/devices/dc/__init__.py
index 404c156e..ece0e166 100644
--- a/midealocal/devices/dc/__init__.py
+++ b/midealocal/devices/dc/__init__.py
@@ -1,5 +1,6 @@
 import logging
 import sys
+from typing import Any
 
 from .message import MessageDCResponse, MessagePower, MessageQuery, MessageStart
 
@@ -34,7 +35,7 @@ def __init__(
         model: str,
         subtype: int,
         customize: str,
-    ):
+    ) -> None:
         super().__init__(
             name=name,
             device_id=device_id,
@@ -55,10 +56,10 @@ def __init__(
             },
         )
 
-    def build_query(self):
+    def build_query(self) -> list[MessageQuery]:
         return [MessageQuery(self._protocol_version)]
 
-    def process_message(self, msg):
+    def process_message(self, msg: bytes) -> dict[str, Any]:
         message = MessageDCResponse(msg)
         _LOGGER.debug(f"[{self.device_id}] Received: {message}")
         new_status = {}
@@ -81,7 +82,8 @@ def process_message(self, msg):
                 new_status[str(status)] = self._attributes[status]
         return new_status
 
-    def set_attribute(self, attr, value):
+    def set_attribute(self, attr: str, value: Any) -> None:
+        message: MessagePower | MessageStart | None = None
         if attr == DeviceAttributes.power:
             message = MessagePower(self._protocol_version)
             message.power = value
diff --git a/midealocal/devices/dc/message.py b/midealocal/devices/dc/message.py
index e20ce50f..b8063514 100644
--- a/midealocal/devices/dc/message.py
+++ b/midealocal/devices/dc/message.py
@@ -7,7 +7,9 @@
 
 
 class MessageDCBase(MessageRequest):
-    def __init__(self, protocol_version, message_type, body_type):
+    def __init__(
+        self, protocol_version: int, message_type: int, body_type: int
+    ) -> None:
         super().__init__(
             device_type=0xDC,
             protocol_version=protocol_version,
@@ -16,12 +18,12 @@ def __init__(self, protocol_version, message_type, body_type):
         )
 
     @property
-    def _body(self):
+    def _body(self) -> bytearray:
         raise NotImplementedError
 
 
 class MessageQuery(MessageDCBase):
-    def __init__(self, protocol_version):
+    def __init__(self, protocol_version: int) -> None:
         super().__init__(
             protocol_version=protocol_version,
             message_type=MessageType.query,
@@ -29,12 +31,12 @@ def __init__(self, protocol_version):
         )
 
     @property
-    def _body(self):
+    def _body(self) -> bytearray:
         return bytearray([])
 
 
 class MessagePower(MessageDCBase):
-    def __init__(self, protocol_version):
+    def __init__(self, protocol_version: int) -> None:
         super().__init__(
             protocol_version=protocol_version,
             message_type=MessageType.set,
@@ -43,13 +45,13 @@ def __init__(self, protocol_version):
         self.power = False
 
     @property
-    def _body(self):
+    def _body(self) -> bytearray:
         power = 0x01 if self.power else 0x00
         return bytearray([power, 0xFF])
 
 
 class MessageStart(MessageDCBase):
-    def __init__(self, protocol_version):
+    def __init__(self, protocol_version: int) -> None:
         super().__init__(
             protocol_version=protocol_version,
             message_type=MessageType.set,
@@ -59,7 +61,7 @@ def __init__(self, protocol_version):
         self.washing_data = bytearray([])
 
     @property
-    def _body(self):
+    def _body(self) -> bytearray:
         if self.start:
             return bytearray([0xFF, 0x01]) + self.washing_data
         else:
@@ -68,25 +70,24 @@ def _body(self):
 
 
 class DCGeneralMessageBody(MessageBody):
-    def __init__(self, body):
+    def __init__(self, body: bytearray) -> None:
         super().__init__(body)
         self.power = body[1] > 0
         self.start = True if body[2] in [2, 6] else False
         self.washing_data = body[3:15]
         self.progress = 0
+        self.time_remaining: float | None = None
         for i in range(7):
             if (body[16] & (1 << i)) > 0:
                 self.progress = i + 1
                 break
         if self.power:
             self.time_remaining = body[17] + body[18] * 60
-        else:
-            self.time_remaining = None
 
 
 class MessageDCResponse(MessageResponse):
-    def __init__(self, message):
-        super().__init__(message)
+    def __init__(self, message: bytes) -> None:
+        super().__init__(bytearray(message))
         if self.message_type in [MessageType.query, MessageType.set] or (
             self.message_type == MessageType.notify1 and self.body_type == 0x04
         ):