chore: refactor to object inerface

Author: joshuagruenstein

adam_modbus/interface.py

@@ -2,7 +2,8 @@
 import socket
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from typing import Literal, assert_never
+from dataclasses import dataclass
+from typing import Literal
 
 
 class AdamConnectionError(RuntimeError):
@@ -13,90 +14,110 @@ class AdamConnectionError(RuntimeError):
 ADAM_CONNECTION_TIMEOUT = 0.1
 
 
-@asynccontextmanager
-async def adam_socket_context(
-    ip: str,
-    port: int = DEFAULT_ADAM_PORT,
-) -> AsyncIterator[socket.socket]:
-    adam_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+@dataclass
+class AdamConnection:
+    socket: socket.socket
+    ip: str
+    port: int
+    timeout: float = ADAM_CONNECTION_TIMEOUT
+    model: str | None = None
 
-    loop = asyncio.get_running_loop()
+    async def _send_and_receive(self, message: str) -> str:
+        loop = asyncio.get_running_loop()
 
-    # according to chatgpt, asyncio does not support timeouts on UDP sockets
-    # so we manually do this with asyncio.wait_for, otherwise it will hang forever
-    adam_sock.setblocking(False)
+        try:
+            await asyncio.wait_for(
+                loop.sock_sendall(self.socket, message.encode("ascii")),
+                self.timeout,
+            )
+            adam_out = await asyncio.wait_for(
+                loop.sock_recv(self.socket, 100), self.timeout
+            )
+        except asyncio.TimeoutError:
+            raise AdamConnectionError("ADAM connection timed out")
 
-    try:
-        await loop.sock_connect(adam_sock, (ip, port))
-        yield adam_sock
+        response = adam_out.decode().strip()
+        return response
 
-    except OSError:
-        raise AdamConnectionError(f"Could not connect to ADAM at {ip}")
+    async def set_digital_out(
+        self,
+        pin: int,
+        value: bool,
+        model: None | Literal["6052"] | Literal["6317"] = None,
+    ) -> None:
+        if model is not None:
+            self.model = model
 
-    finally:
-        adam_sock.close()
+        assert self.model is not None, "Model must be set before setting digital out"
 
+        if self.model == "6052":
+            command = f"#011{pin:x}0{int(value)}\r"
+        elif self.model == "6317":
+            command = f"#01D0{pin:x}{int(value)}\r"
+        else:
+            raise NotImplementedError(
+                f"Digital out not implemented for Adam-{self.model}"
+            )
 
-async def _adam_send_and_receive(message: str, adam_socket: socket.socket) -> str:
-    loop = asyncio.get_running_loop()
+        response = await self._send_and_receive(command)
+        assert response[:3] == ">01", response[:3]
 
-    try:
-        await asyncio.wait_for(
-            loop.sock_sendall(adam_socket, message.encode("ascii")),
-            ADAM_CONNECTION_TIMEOUT,
-        )
-        adam_out = await asyncio.wait_for(
-            loop.sock_recv(adam_socket, 100), ADAM_CONNECTION_TIMEOUT
-        )
-    except asyncio.TimeoutError:
-        raise AdamConnectionError("ADAM connection timed out")
+    async def get_adam_digital_inputs(self) -> list[bool]:
+        response = await self._send_and_receive("$016\r")
+        assert response[:3] == "!01", f"Unexpected response: {response}"
 
-    response = adam_out.decode().strip()
-    return response
+        binary_string = "".join(f"{int(char, 16):0>4b}" for char in response[3:])
+        return [char == "1" for char in binary_string][::-1]
 
+    async def get_adam_analog_inputs(self) -> list[float]:
+        response = await self._send_and_receive("#01\r")
 
-async def set_adam_digital_out(
-    socket: socket.socket,
-    model: Literal["6052"] | Literal["6317"],
-    pin: int,
-    value: bool,
-) -> None:
-    if model == "6052":
-        command = f"#011{pin:x}0{int(value)}\r"
-    elif model == "6317":
-        command = f"#01D0{pin:x}{int(value)}\r"
-    else:
-        assert_never(model)
+        assert response[:3] == ">01", response
+        response_data = response[3:]
 
-    response = await _adam_send_and_receive(command, socket)
-    assert response[:3] == ">01", response[:3]
+        # 7 characters per channel: +00.011
+        assert len(response_data) % 7 == 0, response_data
 
+        return [
+            float(response_data[i * 7 : i * 7 + 7])
+            for i in range(len(response_data) // 7)
+        ]
 
-async def get_adam_digital_inputs(socket: socket.socket) -> list[bool]:
-    response = await _adam_send_and_receive("$016\r", socket)
-    assert response[:3] == "!01", f"Unexpected response: {response}"
+    async def get_adam_model(self) -> str:
+        response = await self._send_and_receive("$01M\r")
 
-    binary_string = "".join(f"{int(char, 16):0>4b}" for char in response[3:])
-    return [char == "1" for char in binary_string][::-1]
+        assert response[:3] == "!01", f"Unexpected response: {response}"
 
+        self.model = response[3:]
 
-async def get_adam_analog_inputs(socket: socket.socket) -> list[float]:
-    response = await _adam_send_and_receive("#01\r", socket)
+        return self.model
 
-    assert response[:3] == ">01", response
-    response_data = response[3:]
 
-    # 7 characters per channel: +00.011
-    assert len(response_data) % 7 == 0, response_data
+@asynccontextmanager
+async def adam_connection_context(
+    ip: str,
+    port: int = DEFAULT_ADAM_PORT,
+    timeout: float = ADAM_CONNECTION_TIMEOUT,
+) -> AsyncIterator[AdamConnection]:
+    adam_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 
-    return [
-        float(response_data[i * 7 : i * 7 + 7]) for i in range(len(response_data) // 7)
-    ]
+    loop = asyncio.get_running_loop()
 
+    # according to chatgpt, asyncio does not support timeouts on UDP sockets
+    # so we manually do this with asyncio.wait_for, otherwise it will hang forever
+    adam_sock.setblocking(False)
 
-async def get_adam_model(socket: socket.socket) -> str:
-    response = await _adam_send_and_receive("$01M\r", socket)
+    try:
+        await loop.sock_connect(adam_sock, (ip, port))
+        yield AdamConnection(
+            adam_sock,
+            ip,
+            port,
+            timeout,
+        )
 
-    assert response[:3] == "!01", f"Unexpected response: {response}"
+    except OSError:
+        raise AdamConnectionError(f"Could not connect to ADAM at {ip}")
 
-    return response[3:]
+    finally:
+        adam_sock.close()