sr-robot3 2024.0.1__py3-none-any.whl → 2025.0.1__py3-none-any.whl

sr/robot3/kch.py

@@ -2,17 +2,28 @@
 from __future__ import annotations
 
 import atexit
+import logging
 import warnings
 from enum import IntEnum, unique
 from typing import Optional
 
+from .exceptions import BoardDisconnectionError, IncorrectBoardError
+from .serial_wrapper import SerialWrapper
+from .utils import IN_SIMULATOR, Board, BoardIdentity, get_simulator_boards
+
 try:
     import RPi.GPIO as GPIO  # isort: ignore
-    HAS_HAT = True
+    HAS_HAT = True if not IN_SIMULATOR else False
 except ImportError:
     HAS_HAT = False
 
 
+logger = logging.getLogger(__name__)
+
+# Only used in the simulator
+BAUDRATE = 115200
+
+
 @unique
 class RobotLEDs(IntEnum):
     """Mapping of LEDs to GPIO Pins."""
@@ -71,6 +82,8 @@ class KCH:
     __slots__ = ('_leds', '_pwm')
 
     def __init__(self) -> None:
+        self._leds: tuple[LED, ...]
+
         if HAS_HAT:
             GPIO.setmode(GPIO.BCM)
             with warnings.catch_warnings():
@@ -81,11 +94,7 @@ class KCH:
                 # suppress this as we know the reason behind the warning
                 GPIO.setup(RobotLEDs.all_leds(), GPIO.OUT, initial=GPIO.LOW)
             self._pwm: Optional[GPIO.PWM] = None
-        self._leds = tuple(
-            LED(rgb_led) for rgb_led in RobotLEDs.user_leds()
-        )
 
-        if HAS_HAT:
             # We are not running cleanup so the LED state persists after the code completes,
             # this will cause a warning with `GPIO.setup()` which we can suppress
             # atexit.register(GPIO.cleanup)
@@ -94,6 +103,25 @@ class KCH:
             # Mypy isn't aware of the version of atexit.register(func, *args)
             atexit.register(GPIO.cleanup, RobotLEDs.START)  # type: ignore[call-arg]
 
+            self._leds = tuple(
+                PhysicalLED(rgb_led) for rgb_led in RobotLEDs.user_leds()
+            )
+        elif IN_SIMULATOR:
+            led_server = LedServer.initialise()
+            if led_server is not None:
+                self._leds = tuple(
+                    SimulationLED(v, led_server)
+                    for v in range(len(RobotLEDs.user_leds()))
+                )
+            else:
+                self._leds = tuple(
+                    LED() for _ in RobotLEDs.user_leds()
+                )
+        else:
+            self._leds = tuple(
+                LED() for _ in RobotLEDs.user_leds()
+            )
+
     @property
     def _start(self) -> bool:
         """Get the state of the start LED."""
@@ -122,7 +150,60 @@ class KCH:
 
 
 class LED:
-    """User programmable LED."""
+    """
+    User programmable LED.
+
+    This is a dummy class to handle the case where this is run on neither the
+    Raspberry Pi nor the simulator.
+    As such, this class does nothing.
+    """
+    __slots__ = ('_led',)
+
+    @property
+    def r(self) -> bool:
+        """Get the state of the Red LED segment."""
+        return False
+
+    @r.setter
+    def r(self, value: bool) -> None:
+        """Set the state of the Red LED segment."""
+
+    @property
+    def g(self) -> bool:
+        """Get the state of the Green LED segment."""
+        return False
+
+    @g.setter
+    def g(self, value: bool) -> None:
+        """Set the state of the Green LED segment."""
+
+    @property
+    def b(self) -> bool:
+        """Get the state of the Blue LED segment."""
+        return False
+
+    @b.setter
+    def b(self, value: bool) -> None:
+        """Set the state of the Blue LED segment."""
+
+    @property
+    def colour(self) -> tuple[bool, bool, bool]:
+        """Get the colour of the user LED."""
+        return False, False, False
+
+    @colour.setter
+    def colour(self, value: tuple[bool, bool, bool]) -> None:
+        """Set the colour of the user LED."""
+        if not isinstance(value, (tuple, list)) or len(value) != 3:
+            raise ValueError("The LED requires 3 values for it's colour")
+
+
+class PhysicalLED(LED):
+    """
+    User programmable LED.
+
+    Used when running on the Raspberry Pi to control the actual LEDs.
+    """
     __slots__ = ('_led',)
 
     def __init__(self, led: tuple[int, int, int]):
@@ -131,47 +212,41 @@ class LED:
     @property
     def r(self) -> bool:
         """Get the state of the Red LED segment."""
-        return GPIO.input(self._led[0]) if HAS_HAT else False
+        return GPIO.input(self._led[0])
 
     @r.setter
     def r(self, value: bool) -> None:
         """Set the state of the Red LED segment."""
-        if HAS_HAT:
-            GPIO.output(self._led[0], GPIO.HIGH if value else GPIO.LOW)
+        GPIO.output(self._led[0], GPIO.HIGH if value else GPIO.LOW)
 
     @property
     def g(self) -> bool:
         """Get the state of the Green LED segment."""
-        return GPIO.input(self._led[1]) if HAS_HAT else False
+        return GPIO.input(self._led[1])
 
     @g.setter
     def g(self, value: bool) -> None:
         """Set the state of the Green LED segment."""
-        if HAS_HAT:
-            GPIO.output(self._led[1], GPIO.HIGH if value else GPIO.LOW)
+        GPIO.output(self._led[1], GPIO.HIGH if value else GPIO.LOW)
 
     @property
     def b(self) -> bool:
         """Get the state of the Blue LED segment."""
-        return GPIO.input(self._led[2]) if HAS_HAT else False
+        return GPIO.input(self._led[2])
 
     @b.setter
     def b(self, value: bool) -> None:
         """Set the state of the Blue LED segment."""
-        if HAS_HAT:
-            GPIO.output(self._led[2], GPIO.HIGH if value else GPIO.LOW)
+        GPIO.output(self._led[2], GPIO.HIGH if value else GPIO.LOW)
 
     @property
     def colour(self) -> tuple[bool, bool, bool]:
         """Get the colour of the user LED."""
-        if HAS_HAT:
-            return (
-                GPIO.input(self._led[0]),
-                GPIO.input(self._led[1]),
-                GPIO.input(self._led[2]),
-            )
-        else:
-            return False, False, False
+        return (
+            GPIO.input(self._led[0]),
+            GPIO.input(self._led[1]),
+            GPIO.input(self._led[2]),
+        )
 
     @colour.setter
     def colour(self, value: tuple[bool, bool, bool]) -> None:
@@ -179,10 +254,161 @@ class LED:
         if not isinstance(value, (tuple, list)) or len(value) != 3:
             raise ValueError("The LED requires 3 values for it's colour")
 
-        if HAS_HAT:
-            GPIO.output(
-                self._led,
-                tuple(
-                    GPIO.HIGH if v else GPIO.LOW for v in value
-                ),
-            )
+        GPIO.output(
+            self._led,
+            tuple(
+                GPIO.HIGH if v else GPIO.LOW for v in value
+            ),
+        )
+
+
+class LedServer(Board):
+    """
+    LED control over a socket.
+
+    Used when running in the simulator to control the simulated LEDs.
+    """
+
+    @staticmethod
+    def get_board_type() -> str:
+        """
+        Return the type of the board.
+
+        :return: The literal string 'KCHv1B'.
+        """
+        return 'KCHv1B'
+
+    def __init__(
+        self,
+        serial_port: str,
+        initial_identity: BoardIdentity | None = None,
+    ) -> None:
+        if initial_identity is None:
+            initial_identity = BoardIdentity()
+        self._serial = SerialWrapper(
+            serial_port,
+            BAUDRATE,
+            identity=initial_identity,
+        )
+
+        self._identity = self.identify()
+        if self._identity.board_type != self.get_board_type():
+            raise IncorrectBoardError(self._identity.board_type, self.get_board_type())
+        self._serial.set_identity(self._identity)
+
+        # Reset the board to a known state
+        self._serial.write('*RESET')
+
+    @classmethod
+    def initialise(cls) -> 'LedServer' | None:
+        """Initialise the LED server using simulator discovery."""
+        # The filter here is the name of the emulated board in the simulator
+        boards = get_simulator_boards('LedBoard')
+
+        if not boards:
+            return None
+
+        board_info = boards[0]
+
+        # Create board identity from the info given
+        initial_identity = BoardIdentity(
+            manufacturer='sbot_simulator',
+            board_type=board_info.type_str,
+            asset_tag=board_info.serial_number,
+        )
+
+        try:
+            board = cls(board_info.url, initial_identity)
+        except BoardDisconnectionError:
+            logger.warning(
+                f"Simulator specified LED board at port {board_info.url!r}, "
+                "could not be identified. Ignoring this device")
+            return None
+        except IncorrectBoardError as err:
+            logger.warning(
+                f"Board returned type {err.returned_type!r}, "
+                f"expected {err.expected_type!r}. Ignoring this device")
+            return None
+
+        return board
+
+    def identify(self) -> BoardIdentity:
+        """
+        Get the identity of the board.
+
+        :return: The identity of the board.
+        """
+        response = self._serial.query('*IDN?')
+        return BoardIdentity(*response.split(':'))
+
+    def set_leds(self, led_num: int, value: tuple[bool, bool, bool]) -> None:
+        """Set the colour of the LED."""
+        self._serial.write(f'LED:{led_num}:SET:{value[0]:d}:{value[1]:d}:{value[2]:d}')
+
+    def get_leds(self, led_num: int) -> tuple[bool, bool, bool]:
+        """Get the colour of the LED."""
+        response = self._serial.query(f'LED:{led_num}:GET?')
+        red, green, blue = response.split(':')
+        return bool(int(red)), bool(int(green)), bool(int(blue))
+
+
+class SimulationLED(LED):
+    """
+    User programmable LED.
+
+    Used when running in the simulator to control the simulated LEDs.
+    """
+    __slots__ = ('_led_num', '_server')
+
+    def __init__(self, led_num: int, server: LedServer) -> None:
+        self._led_num = led_num
+        self._server = server
+
+    @property
+    def r(self) -> bool:
+        """Get the state of the Red LED segment."""
+        return self._server.get_leds(self._led_num)[0]
+
+    @r.setter
+    def r(self, value: bool) -> None:
+        """Set the state of the Red LED segment."""
+        # Fetch the current state of the LED so we can update only the red value
+        current = self._server.get_leds(self._led_num)
+        self._server.set_leds(self._led_num, (value, current[1], current[2]))
+
+    @property
+    def g(self) -> bool:
+        """Get the state of the Green LED segment."""
+        return self._server.get_leds(self._led_num)[1]
+
+    @g.setter
+    def g(self, value: bool) -> None:
+        """Set the state of the Green LED segment."""
+        # Fetch the current state of the LED so we can update only the green value
+        current = self._server.get_leds(self._led_num)
+        self._server.set_leds(self._led_num, (current[0], value, current[2]))
+
+    @property
+    def b(self) -> bool:
+        """Get the state of the Blue LED segment."""
+        return self._server.get_leds(self._led_num)[2]
+
+    @b.setter
+    def b(self, value: bool) -> None:
+        """Set the state of the Blue LED segment."""
+        # Fetch the current state of the LED so we can update only the blue value
+        current = self._server.get_leds(self._led_num)
+        self._server.set_leds(self._led_num, (current[0], current[1], value))
+
+    @property
+    def colour(self) -> tuple[bool, bool, bool]:
+        """Get the colour of the user LED."""
+        return self._server.get_leds(self._led_num)
+
+    @colour.setter
+    def colour(self, value: tuple[bool, bool, bool]) -> None:
+        """Set the colour of the user LED."""
+        if not isinstance(value, (tuple, list)) or len(value) != 3:
+            raise ValueError("The LED requires 3 values for it's colour")
+
+        self._server.set_leds(self._led_num, value)

sr/robot3/motor_board.py

@@ -13,8 +13,8 @@ from .exceptions import IncorrectBoardError
 from .logging import log_to_debug
 from .serial_wrapper import SerialWrapper
 from .utils import (
-    Board, BoardIdentity, float_bounds_check,
-    get_USB_identity, map_to_float, map_to_int,
+    IN_SIMULATOR, Board, BoardIdentity, float_bounds_check,
+    get_simulator_boards, get_USB_identity, map_to_float, map_to_int,
 )
 
 logger = logging.getLogger(__name__)
@@ -113,17 +113,46 @@ class MotorBoard(Board):
                 f"expected {err.expected_type!r}. Ignoring this device")
             return None
         except Exception:
-            if initial_identity is not None and initial_identity.board_type == 'manual':
-                logger.warning(
-                    f"Manually specified motor board at port {serial_port!r}, "
-                    "could not be identified. Ignoring this device")
-            else:
-                logger.warning(
-                    f"Found motor board-like serial port at {serial_port!r}, "
-                    "but it could not be identified. Ignoring this device")
+            if initial_identity is not None:
+                if initial_identity.board_type == 'manual':
+                    logger.warning(
+                        f"Manually specified motor board at port {serial_port!r}, "
+                        "could not be identified. Ignoring this device")
+                elif initial_identity.manufacturer == 'sbot_simulator':
+                    logger.warning(
+                        f"Simulator specified motor board at port {serial_port!r}, "
+                        "could not be identified. Ignoring this device")
+                return None
+
+            logger.warning(
+                f"Found motor board-like serial port at {serial_port!r}, "
+                "but it could not be identified. Ignoring this device")
             return None
         return board
 
+    @classmethod
+    def _get_simulator_boards(cls) -> MappingProxyType[str, MotorBoard]:
+        """
+        Get the simulator boards.
+
+        :return: A mapping of board serial numbers to boards.
+        """
+        boards = {}
+        # The filter here is the name of the emulated board in the simulator
+        for board_info in get_simulator_boards('MotorBoard'):
+
+            # Create board identity from the info given
+            initial_identity = BoardIdentity(
+                manufacturer='sbot_simulator',
+                board_type=board_info.type_str,
+                asset_tag=board_info.serial_number,
+            )
+            if (board := cls._get_valid_board(board_info.url, initial_identity)) is None:
+                continue
+
+            boards[board._identity.asset_tag] = board
+        return MappingProxyType(boards)
+
     @classmethod
     def _get_supported_boards(
         cls, manual_boards: Optional[list[str]] = None,
@@ -137,6 +166,9 @@ class MotorBoard(Board):
             to connect to, defaults to None
         :return: A mapping of serial numbers to motor boards.
         """
+        if IN_SIMULATOR:
+            return cls._get_simulator_boards()
+
         boards = {}
         serial_ports = comports()
         for port in serial_ports:

sr/robot3/mqtt.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import atexit
 import json
 import logging
+import time
 from threading import Lock
 from typing import Any, Callable, Optional, TypedDict, Union
 from urllib.parse import urlparse
@@ -17,7 +18,7 @@ class MQTTClient:
         self,
         client_name: Optional[str] = None,
         topic_prefix: Optional[str] = None,
-        mqtt_version: int = mqtt.MQTTv5,
+        mqtt_version: mqtt.MQTTProtocolVersion = mqtt.MQTTProtocolVersion.MQTTv5,
         use_tls: Union[bool, str] = False,
         username: str = '',
         password: str = '',
@@ -29,7 +30,11 @@ class MQTTClient:
         self.topic_prefix = topic_prefix
         self._client_name = client_name
 
-        self._client = mqtt.Client(client_id=client_name, protocol=mqtt_version)
+        self._client = mqtt.Client(
+            callback_api_version=mqtt.CallbackAPIVersion.VERSION2,
+            client_id=client_name,
+            protocol=mqtt_version,
+        )
         self._client.on_connect = self._on_connect
 
         if use_tls:
@@ -143,19 +148,29 @@ class MQTTClient:
         """Wrap a payload up to be decodable as JSON."""
         if isinstance(payload, bytes):
             payload = payload.decode('utf-8')
+
+        payload_dict = {
+            "timestamp": time.time(),
+            "data": payload,
+        }
+
         self.publish(
             topic,
-            json.dumps({"data": payload}),
+            json.dumps(payload_dict),
             retain=retain, abs_topic=abs_topic)
 
     def _on_connect(
-        self, client: mqtt.Client, userdata: Any, flags: dict[str, int], rc: int,
-        properties: Optional[mqtt.Properties] = None,
+        self,
+        client: mqtt.Client,
+        userdata: Any,
+        connect_flags: mqtt.ConnectFlags,
+        reason_code: mqtt.ReasonCode,
+        properties: mqtt.Properties | None = None,
     ) -> None:
         """Callback run each time the client connects to the broker."""
-        if rc != mqtt.CONNACK_ACCEPTED:
+        if reason_code.is_failure:
             LOGGER.warning(
-                f"Failed to connect to MQTT broker. Return code: {mqtt.error_string(rc)}"
+                f"Failed to connect to MQTT broker. Return code: {reason_code.getName()}"  # type: ignore[no-untyped-call] # noqa: E501
             )
             return
 
@@ -182,7 +197,7 @@ def unpack_mqtt_url(url: str) -> MQTTVariables:
     The URL should be in the format:
     mqtt[s]://[<username>[:<password>]]@<host>[:<port>]/<topic_root>
     """
-    url_parts = urlparse(url)
+    url_parts = urlparse(url, allow_fragments=False)
 
     if url_parts.scheme not in ('mqtt', 'mqtts'):
         raise ValueError(f"Invalid scheme: {url_parts.scheme}")

sr/robot3/power_board.py

@@ -6,17 +6,21 @@ import logging
 from enum import IntEnum
 from time import sleep
 from types import MappingProxyType
-from typing import NamedTuple, Optional
+from typing import Callable, NamedTuple, Optional
 
 from serial.tools.list_ports import comports
 
 from .exceptions import IncorrectBoardError
 from .logging import log_to_debug
 from .serial_wrapper import SerialWrapper
-from .utils import Board, BoardIdentity, float_bounds_check, get_USB_identity
+from .utils import (
+    IN_SIMULATOR, Board, BoardIdentity, float_bounds_check,
+    get_simulator_boards, get_USB_identity,
+)
 
 logger = logging.getLogger(__name__)
 BAUDRATE = 115200  # Since the power board is a USB device, this is ignored
+_SLEEP_FN = sleep  # For use in the simulator
 
 
 class PowerOutputPosition(IntEnum):
@@ -132,20 +136,50 @@ class PowerBoard(Board):
                 f"expected {err.expected_type!r}. Ignoring this device")
             return None
         except Exception:
-            if initial_identity is not None and initial_identity.board_type == 'manual':
-                logger.warning(
-                    f"Manually specified power board at port {serial_port!r}, "
-                    "could not be identified. Ignoring this device")
-            else:
-                logger.warning(
-                    f"Found power board-like serial port at {serial_port!r}, "
-                    "but it could not be identified. Ignoring this device")
+            if initial_identity is not None:
+                if initial_identity.board_type == 'manual':
+                    logger.warning(
+                        f"Manually specified power board at port {serial_port!r}, "
+                        "could not be identified. Ignoring this device")
+                elif initial_identity.manufacturer == 'sbot_simulator':
+                    logger.warning(
+                        f"Simulator specified power board at port {serial_port!r}, "
+                        "could not be identified. Ignoring this device")
+
+            logger.warning(
+                f"Found power board-like serial port at {serial_port!r}, "
+                "but it could not be identified. Ignoring this device")
             return None
         return board
 
+    @classmethod
+    def _get_simulator_boards(cls) -> MappingProxyType[str, PowerBoard]:
+        """
+        Get the simulator boards.
+
+        :return: A mapping of board serial numbers to boards.
+        """
+        boards = {}
+        # The filter here is the name of the emulated board in the simulator
+        for board_info in get_simulator_boards('PowerBoard'):
+
+            # Create board identity from the info given
+            initial_identity = BoardIdentity(
+                manufacturer='sbot_simulator',
+                board_type=board_info.type_str,
+                asset_tag=board_info.serial_number,
+            )
+            if (board := cls._get_valid_board(board_info.url, initial_identity)) is None:
+                continue
+
+            boards[board._identity.asset_tag] = board
+        return MappingProxyType(boards)
+
     @classmethod
     def _get_supported_boards(
-        cls, manual_boards: Optional[list[str]] = None,
+        cls,
+        manual_boards: Optional[list[str]] = None,
+        sleep_fn: Optional[Callable[[float], None]] = None,
     ) -> MappingProxyType[str, PowerBoard]:
         """
         Find all connected power boards.
@@ -156,6 +190,12 @@ class PowerBoard(Board):
             to connect to, defaults to None
         :return: A mapping of serial numbers to power boards.
         """
+        if sleep_fn is not None:
+            global _SLEEP_FN
+            _SLEEP_FN = sleep_fn
+        if IN_SIMULATOR:
+            return cls._get_simulator_boards()
+
         boards = {}
         serial_ports = comports()
         for port in serial_ports:
@@ -501,13 +541,13 @@ class Piezo:
             frequency, 8, 10_000, "Frequency must be between 8 and 10000Hz"))
         duration_ms = int(float_bounds_check(
             duration * 1000, 0, 2**31 - 1,
-            f"Duration is a float in seconds between 0 and {(2**31-1)/1000:,.0f}"))
+            f"Duration is a float in seconds between 0 and {(2**31 - 1) / 1000:,.0f}"))
 
         cmd = f'NOTE:{frequency_int}:{duration_ms}'
         self._serial.write(cmd)
 
         if blocking:
-            sleep(duration)
+            _SLEEP_FN(duration)
 
     def __repr__(self) -> str:
         return f"<{self.__class__.__qualname__}: {self._serial}>"