chore: comments

Author: Lash-L

roborock/version_1_apis/roborock_local_client_v1.py

@@ -3,6 +3,7 @@
 from asyncio import Lock, TimerHandle, Transport, get_running_loop
 from collections.abc import Callable
 from dataclasses import dataclass
+from enum import StrEnum
 
 import async_timeout
 
@@ -18,6 +19,13 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+class LocalProtocolVersion(StrEnum):
+    """Supported local protocol versions. Different from vacuum protocol versions."""
+
+    L01 = "L01"
+    V1 = "1.0"
+
+
 @dataclass
 class _LocalProtocol(asyncio.Protocol):
     """Callbacks for the Roborock local client transport."""
@@ -37,7 +45,12 @@ def connection_lost(self, exc: Exception | None) -> None:
 class RoborockLocalClientV1(RoborockClientV1, RoborockClient):
     """Roborock local client for v1 devices."""
 
-    def __init__(self, device_data: DeviceData, queue_timeout: int = 4, version: str | None = None):
+    def __init__(
+        self,
+        device_data: DeviceData,
+        queue_timeout: int = 4,
+        local_protocol_version: LocalProtocolVersion | None = None,
+    ):
         """Initialize the Roborock local client."""
         if device_data.host is None:
             raise RoborockException("Host is required")
@@ -50,13 +63,17 @@ def __init__(self, device_data: DeviceData, queue_timeout: int = 4, version: str
         RoborockClientV1.__init__(self, device_data, security_data=None)
         RoborockClient.__init__(self, device_data)
         self._local_protocol = _LocalProtocol(self._data_received, self._connection_lost)
-        self._version = version
+        self._local_protocol_version = local_protocol_version
         self._connect_nonce = get_next_int(10000, 32767)
         self._ack_nonce: int | None = None
         self._set_encoder_decoder()
         self.queue_timeout = queue_timeout
         self._logger = RoborockLoggerAdapter(device_data.device.name, _LOGGER)
 
+    @property
+    def local_protocol_version(self) -> LocalProtocolVersion:
+        return LocalProtocolVersion.V1 if self._local_protocol_version is None else self._local_protocol_version
+
     def _data_received(self, message):
         """Called when data is received from the transport."""
         parsed_msg = self._decoder(message)
@@ -111,16 +128,21 @@ async def async_disconnect(self) -> None:
             self._sync_disconnect()
 
     def _set_encoder_decoder(self):
-        """Updates the encoder decoder. For L01 these are updated with nonces after the first hello."""
+        """Updates the encoder decoder. These are updated with nonces after the first hello.
+        Only L01 uses the nonces."""
         self._encoder = create_local_encoder(self.device_info.device.local_key, self._connect_nonce, self._ack_nonce)
         self._decoder = create_local_decoder(self.device_info.device.local_key, self._connect_nonce, self._ack_nonce)
 
-    async def _do_hello(self, version: str) -> bool:
+    async def _do_hello(self, local_protocol_version: LocalProtocolVersion) -> bool:
         """Perform the initial handshaking."""
-        self._logger.debug(f"Attempting to use the {version} protocol for client {self.device_info.device.duid}...")
+        self._logger.debug(
+            "Attempting to use the %s protocol for client %s...",
+            local_protocol_version,
+            self.device_info.device.duid,
+        )
         request = RoborockMessage(
             protocol=RoborockMessageProtocol.HELLO_REQUEST,
-            version=version.encode(),
+            version=local_protocol_version.encode(),
             random=self._connect_nonce,
             seq=1,
         )
@@ -132,31 +154,39 @@ async def _do_hello(self, version: str) -> bool:
             )
             self._ack_nonce = response.random
             self._set_encoder_decoder()
-            self._version = version
-            self._logger.debug(f"Client {self.device_info.device.duid} speaks the {version} protocol.")
+            self._local_protocol_version = local_protocol_version
+
+            self._logger.debug(
+                "Client %s speaks the %s protocol.",
+                self.device_info.device.duid,
+                local_protocol_version,
+            )
             return True
         except RoborockException as e:
             self._logger.debug(
-                f"Client {self.device_info.device.duid} did not respond or does not speak the {version} protocol. {e}"
+                "Client %s did not respond or does not speak the %s protocol. %s",
+                self.device_info.device.duid,
+                local_protocol_version,
+                e,
             )
             return False
 
     async def hello(self):
         """Send hello to the device to negotiate protocol."""
-        if self._version:
+        if self._local_protocol_version:
             # version is forced
-            if not await self._do_hello(self._version):
-                raise RoborockException(f"Failed to connect to device with protocol {self._version}")
+            if not await self._do_hello(self._local_protocol_version):
+                raise RoborockException(f"Failed to connect to device with protocol {self._local_protocol_version}")
         else:
             # try 1.0, then L01
-            if not await self._do_hello("1.0"):
-                if not await self._do_hello("L01"):
+            if not await self._do_hello(LocalProtocolVersion.V1):
+                if not await self._do_hello(LocalProtocolVersion.L01):
                     raise RoborockException("Failed to connect to device with any known protocol")
 
     async def ping(self) -> None:
-        # Realistically, this should be set here, but this is to be safe and for typing.
-        version = b"1.0" if self._version is None else self._version.encode()
-        ping_message = RoborockMessage(protocol=RoborockMessageProtocol.PING_REQUEST, version=version)
+        ping_message = RoborockMessage(
+            protocol=RoborockMessageProtocol.PING_REQUEST, version=self.local_protocol_version.encode()
+        )
         await self._send_message(
             roborock_message=ping_message,
             request_id=ping_message.seq,
@@ -180,7 +210,8 @@ async def _send_command(
             raise RoborockException(f"Method {method} is not supported over local connection")
         request_message = RequestMessage(method=method, params=params)
         roborock_message = request_message.encode_message(
-            RoborockMessageProtocol.GENERAL_REQUEST, version=self._version if self._version is not None else "1.0"
+            RoborockMessageProtocol.GENERAL_REQUEST,
+            version=self.local_protocol_version,
         )
         self._logger.debug("Building message id %s for method %s", request_message.request_id, method)
         return await self._send_message(