feat: enhance type hinting across client and control modules for improved clarity

Author: JakubAndrysek

src/wokwi_client/client.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: MIT
 
+import typing
 from pathlib import Path
 from typing import Any, Optional
 
@@ -181,7 +182,7 @@ async def serial_monitor_cat(self) -> None:
         async for line in monitor_lines(self._transport):
             print(line.decode("utf-8"), end="", flush=True)
 
-    async def serial_write(self, data: bytes | str | list[int]) -> None:
+    async def serial_write(self, data: typing.Union[bytes, str, list[int]]) -> None:
         """Write data to the simulation serial monitor interface."""
         await write_serial(self._transport, data)
 
@@ -210,7 +211,9 @@ async def listen_pin(self, part: str, pin: str, listen: bool = True) -> Response
         """
         return await pin_listen(self._transport, part=part, pin=pin, listen=listen)
 
-    async def set_control(self, part: str, control: str, value: int | bool | float) -> ResponseMessage:
+    async def set_control(
+        self, part: str, control: str, value: typing.Union[int, bool, float]
+    ) -> ResponseMessage:
         """Set a control value (e.g. simulate button press).
 
         Args:

src/wokwi_client/client_sync.py

@@ -2,58 +2,67 @@
 #
 # SPDX-License-Identifier: MIT
 
+from __future__ import annotations
+
 import asyncio
 import logging
 import threading
 import typing as t
+from concurrent.futures import Future
 from pathlib import Path
 
 from wokwi_client import WokwiClient
 from wokwi_client.serial import monitor_lines as monitor_serial_lines
 
+if t.TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 class WokwiClientSync:
     """Synchronous wrapper around the async WokwiClient."""
 
-    def __init__(self, token: str, server: t.Optional[str] = None):
+    token: str
+    server: t.Optional[str]
+    _loop: t.Optional[asyncio.AbstractEventLoop]
+    _loop_thread: t.Optional[threading.Thread]
+    _client: t.Optional[WokwiClient]
+    _monitor_task: t.Optional[Future[t.Any]]
+    _connected: bool
+
+    def __init__(self, token: str, server: t.Optional[str] = None) -> None:
         self.token = token
         self.server = server
         self._loop = None
         self._loop_thread = None
         self._client = None
+        self._monitor_task = None
         self._connected = False
 
-    def _ensure_loop(self):
-        """Ensure the async event loop is running."""
+    def _ensure_loop(self) -> None:
         if self._loop is None:
             self._loop = asyncio.new_event_loop()
             self._loop_thread = threading.Thread(target=self._loop.run_forever, daemon=True)
             self._loop_thread.start()
 
-    def _run_async(self, coro, timeout=30):
-        """Run an async coroutine synchronously."""
+    def _run_async(self, coro: t.Coroutine[t.Any, t.Any, t.Any], timeout: float = 30) -> t.Any:
         self._ensure_loop()
+        assert self._loop is not None
         future = asyncio.run_coroutine_threadsafe(coro, self._loop)
         return future.result(timeout=timeout)
 
-    def connect(self):
-        """Connect to Wokwi server."""
+    def connect(self) -> t.Dict[str, t.Any]:
         if not self._connected:
             self._client = WokwiClient(self.token, self.server)
-            result = self._run_async(self._client.connect())
+            result: t.Dict[str, t.Any] = t.cast(t.Dict[str, t.Any], self._run_async(self._client.connect()))
             self._connected = True
             return result
         return {}
 
-    def disconnect(self):
-        """Disconnect from Wokwi server."""
+    def disconnect(self) -> None:
         if self._connected and self._client:
             try:
-                # Stop any ongoing monitor task
-                if hasattr(self, '_monitor_task') and self._monitor_task:
+                if self._monitor_task:
                     self._monitor_task.cancel()
-
-                # Disconnect the client
                 self._run_async(self._client.disconnect(), timeout=5)
             except Exception as e:
                 logging.debug(f"Error during disconnect: {e}")
@@ -70,79 +79,85 @@ def disconnect(self):
                         self._loop = None
                         self._loop_thread = None
 
-    def upload(self, name: str, content: bytes):
-        """Upload a file to the simulator from bytes content."""
+    def upload(self, name: str, content: bytes) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.upload(name, content))
 
-    def upload_file(self, filename: str, local_path: t.Optional[Path] = None):
-        """Upload a file to the simulator."""
+    def upload_file(self, filename: str, local_path: t.Optional[Path] = None) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.upload_file(filename, local_path))
 
-    def start_simulation(self, firmware: str, elf: t.Optional[str] = None, pause: bool = False, chips: list[str] = []):
-        """Start a simulation."""
+    def start_simulation(
+        self,
+        firmware: str,
+        elf: t.Optional[str] = None,
+        pause: bool = False,
+        chips: t.Optional[t.List[str]] = None,
+    ) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
-        return self._run_async(self._client.start_simulation(firmware, elf, pause, chips))
+        assert self._client is not None
+        return self._run_async(self._client.start_simulation(firmware, elf, pause, chips or []))
 
-    def pause_simulation(self):
-        """Pause the running simulation."""
+    def pause_simulation(self) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.pause_simulation())
 
-    def resume_simulation(self, pause_after: t.Optional[int] = None):
-        """Resume the simulation, optionally pausing after a given number of nanoseconds."""
+    def resume_simulation(self, pause_after: t.Optional[int] = None) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.resume_simulation(pause_after))
 
-    def wait_until_simulation_time(self, seconds: float):
-        """Pause and resume the simulation until the given simulation time (in seconds) is reached."""
+    def wait_until_simulation_time(self, seconds: float) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.wait_until_simulation_time(seconds))
 
-    def restart_simulation(self, pause: bool = False):
-        """Restart the simulation, optionally starting paused."""
+    def restart_simulation(self, pause: bool = False) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.restart_simulation(pause))
 
-    def serial_monitor_cat(self):
-        """Print serial monitor output to stdout as it is received from the simulation."""
+    def serial_monitor_cat(self) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.serial_monitor_cat())
 
-    def write_serial(self, data: t.Union[bytes, str, list[int]]):
-        """Write data to serial."""
+    def write_serial(self, data: t.Union[bytes, str, t.List[int]]) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.serial_write(data))
 
-    def read_pin(self, part: str, pin: str):
-        """Read the current state of a pin."""
+    def read_pin(self, part: str, pin: str) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.read_pin(part, pin))
 
-    def listen_pin(self, part: str, pin: str, listen: bool = True):
-        """Start or stop listening for changes on a pin."""
+    def listen_pin(self, part: str, pin: str, listen: bool = True) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.listen_pin(part, pin, listen))
 
-    def monitor_serial(self, callback):
-        """Start monitoring serial output with a callback."""
+    def monitor_serial(self, callback: t.Callable[[bytes], None]) -> None:
         if not self._connected:
             raise RuntimeError("Client not connected")
 
-        async def _monitor():
+        async def _monitor() -> None:
             try:
+                assert self._client is not None
                 async for line in monitor_serial_lines(self._client._transport):
                     if not self._connected:
                         break
@@ -154,27 +169,24 @@ async def _monitor():
             except Exception as e:
                 logging.error(f"Error in serial monitor: {e}")
 
-        # Start monitoring in background
+        assert self._loop is not None
         self._monitor_task = asyncio.run_coroutine_threadsafe(_monitor(), self._loop)
 
-    def set_control(self, part: str, control: str, value: t.Union[int, bool, float]):
-        """Set control value for a part."""
+    def set_control(self, part: str, control: str, value: t.Union[int, bool, float]) -> t.Any:
         if not self._connected:
             raise RuntimeError("Client not connected")
+        assert self._client is not None
         return self._run_async(self._client.set_control(part, control, value))
 
     @property
-    def version(self):
-        """Get client version."""
+    def version(self) -> str:
         if self._client:
             return self._client.version
-        # Return a default version if client not initialized yet
         client = WokwiClient(self.token, self.server)
         return client.version
 
     @property
-    def last_pause_nanos(self):
-        """Get the last pause time in nanoseconds."""
+    def last_pause_nanos(self) -> int:
         if self._client:
             return self._client.last_pause_nanos
         return 0

src/wokwi_client/control.py

@@ -11,12 +11,14 @@
 #
 # SPDX-License-Identifier: MIT
 
+import typing
+
 from .protocol_types import ResponseMessage
 from .transport import Transport
 
 
 async def set_control(
-    transport: Transport, *, part: str, control: str, value: int | bool | float
+    transport: Transport, *, part: str, control: str, value: typing.Union[int, bool, float]
 ) -> ResponseMessage:
     """Set a control value on a part (e.g. simulate button press/release).
 

src/wokwi_client/pins.py

@@ -15,9 +15,7 @@
 from .transport import Transport
 
 
-async def pin_read(
-    transport: Transport, *, part: str, pin: str
-) -> ResponseMessage:
+async def pin_read(transport: Transport, *, part: str, pin: str) -> ResponseMessage:
     """Read the state of a pin.
 
     Args:
@@ -44,6 +42,4 @@ async def pin_listen(
         listen: True to start listening, False to stop.
     """
 
-    return await transport.request(
-        "pin:listen", {"part": part, "pin": pin, "listen": listen}
-    )
+    return await transport.request("pin:listen", {"part": part, "pin": pin, "listen": listen})

src/wokwi_client/serial.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: MIT
 
+import typing
 from collections.abc import AsyncGenerator, Iterable
 
 from .event_queue import EventQueue
@@ -16,7 +17,7 @@ async def monitor_lines(transport: Transport) -> AsyncGenerator[bytes, None]:
             yield bytes(event_msg["payload"]["bytes"])
 
 
-async def write_serial(transport: Transport, data: bytes | str | Iterable[int]) -> None:
+async def write_serial(transport: Transport, data: typing.Union[bytes, str, Iterable[int]]) -> None:
     """Write data to the serial monitor.
 
     Accepts bytes, str (encoded as utf-8), or an iterable of integer byte values.

tests/utils.py

@@ -11,9 +11,12 @@
 import subprocess
 import sys
 from collections.abc import Mapping
+from subprocess import CompletedProcess
 
 
-def run_example_module(module: str, *, sleep_time: str = "1", extra_env: Mapping[str, str] | None = None) -> subprocess.CompletedProcess:
+def run_example_module(
+    module: str, *, sleep_time: str = "1", extra_env: Mapping[str, str] | None = None
+) -> CompletedProcess[str]:
     """Run an example module with a short simulation time.
 
     Requires WOKWI_CLI_TOKEN to be set in the environment.