sr-robot3 2024.0.0rc2__py3-none-any.whl → 2025.0.0__py3-none-any.whl

sr/robot3/raw_serial.py

@@ -2,13 +2,15 @@
 from __future__ import annotations
 
 import logging
+from contextlib import contextmanager
 from time import sleep
 from types import MappingProxyType
-from typing import NamedTuple, Optional
+from typing import Generator, NamedTuple, Optional
 
-from serial import Serial, serial_for_url
+from serial import Serial, SerialException, serial_for_url
 from serial.tools.list_ports import comports
 
+from .exceptions import BoardDisconnectionError
 from .logging import log_to_debug
 from .utils import Board, BoardIdentity, get_USB_identity
 
@@ -33,7 +35,7 @@ class RawSerial(Board):
     :param baudrate: The baudrate to use when connecting.
     :param identity: The identity of the board, as reported by the USB descriptor.
     """
-    __slots__ = ('_serial', '_identity')
+    __slots__ = ('_serial', '_identity', 'delay_after_connect')
 
     @staticmethod
     def get_board_type() -> str:
@@ -63,9 +65,13 @@ class RawSerial(Board):
             serial_port,
             baudrate=baudrate,
             timeout=0.5,  # 500ms timeout
+            do_not_open=True,
         )
 
-        sleep(2)  # Wait for boards to reset after connecting
+        # Certain boards will reset when the serial port is opened
+        # Wait for the board to be ready to receive data
+        self.delay_after_connect = 2
+        self._reconnect()
 
     @classmethod
     def _get_supported_boards(
@@ -75,7 +81,7 @@ class RawSerial(Board):
         """
         Discover connected serial devices filtered by serial number.
 
-        :param serial_ports: A list of serial ports to check (serial_number, baudrate),
+        :param serial_devices: A list of serial ports to check (serial_number, baudrate),
             these are matched by serial number.
         :return: A mapping of board serial numbers to Arduinos
         """
@@ -123,9 +129,9 @@ class RawSerial(Board):
 
     @property
     @log_to_debug
-    def raw(self) -> Serial:
+    def port(self) -> Serial:
         """
-        The raw serial port.
+        The pySerial Serial object.
 
         :return: The raw serial port.
         """
@@ -138,7 +144,8 @@ class RawSerial(Board):
 
         :param data: The bytes to send.
         """
-        self._serial.write(data)
+        with self._handle_serial_reconnect():
+            self._serial.write(data)
 
     @log_to_debug
     def read(self, size: int) -> bytes:
@@ -148,7 +155,8 @@ class RawSerial(Board):
         :param size: The number of bytes to read.
         :return: The line read from the port.
         """
-        return self._serial.read(size)
+        with self._handle_serial_reconnect():
+            return self._serial.read(size)
 
     @log_to_debug
     def read_until(self, terminator: bytes = b'\n') -> bytes:
@@ -158,7 +166,60 @@ class RawSerial(Board):
         :param terminator: The terminator character to stop reading at.
         :return: The data read from the port.
         """
-        return self._serial.read_until(terminator)
+        with self._handle_serial_reconnect():
+            return self._serial.read_until(terminator)
+
+    @contextmanager
+    def _handle_serial_reconnect(self) -> Generator[None, None, None]:
+        """
+        A context manager to handle reconnecting the serial ports on faults.
+        """
+        if not self._serial.is_open:
+            if not self._reconnect():
+                # If the serial port cannot be opened raise an error
+                raise BoardDisconnectionError(
+                    f'Connection to port {self._identity.board_type}:'
+                    f'{self._identity.asset_tag} could not be established',
+                )
+        try:
+            yield
+        except (SerialException, OSError):
+            # Serial connection failed, close the port and raise an error
+            self._serial.close()
+            logger.warning(
+                f'Port {self._identity.board_type}:{self._identity.asset_tag} disconnected',
+            )
+            raise BoardDisconnectionError(
+                f'Port {self._identity.board_type}:{self._identity.asset_tag} '
+                'disconnected during transaction',
+            )
+
+    def _reconnect(self) -> bool:
+        """
+        Connect to the class's serial port.
+
+        This is called automatically when a message is sent to the board if the port is closed.
+
+        :raises serial.SerialException: If the serial port cannot be opened.
+        :return: True if the serial port was opened successfully, False otherwise.
+        """
+        try:
+            self._serial.open()
+            # Wait for the board to be ready to receive data
+            # Certain boards will reset when the serial port is opened
+            sleep(self.delay_after_connect)
+            self._serial.reset_input_buffer()
+        except SerialException:
+            logger.error(
+                'Failed to connect to port '
+                f'{self._identity.board_type}:{self._identity.asset_tag}',
+            )
+            return False
+
+        logger.info(
+            f'Connected to port {self._identity.board_type}: {self._identity.asset_tag}',
+        )
+        return True
 
     def __repr__(self) -> str:
         return f"<{self.__class__.__qualname__}: {self._serial}>"

sr/robot3/robot.py

@@ -3,8 +3,10 @@ from __future__ import annotations
 
 import itertools
 import logging
+import os
 import time
 from pathlib import Path
+from socket import socket
 from types import MappingProxyType
 from typing import Mapping, Optional
 
@@ -20,7 +22,8 @@ from .motor_board import MotorBoard
 from .power_board import Note, PowerBoard
 from .raw_serial import RawSerial
 from .servo_board import ServoBoard
-from .utils import ensure_atexit_on_term, obtain_lock, singular
+from .simulator.time_server import TimeServer
+from .utils import IN_SIMULATOR, ensure_atexit_on_term, obtain_lock, singular
 
 logger = logging.getLogger(__name__)
 
@@ -36,13 +39,16 @@ class Robot:
     :param wait_for_start: Wait in the constructor until the start button is pressed,
         defaults to True
     :param trace_logging: Enable trace level logging to the console, defaults to False
+    :param ignored_arduinos: A list of Arduino serial numbers to avoid connecting to
     :param manual_boards: A dictionary of board types to a list of serial port paths
         to allow for connecting to boards that are not automatically detected, defaults to None
     :param raw_ports: A list of serial number, baudrate tuples to try connecting to.
+    :param no_powerboard: If True, initialize the robot without a powerboard, defaults to False
     """
     __slots__ = (
         '_lock', '_metadata', '_power_board', '_motor_boards', '_servo_boards',
         '_arduinos', '_cameras', '_mqtt', '_astoria', '_kch', '_raw_ports',
+        '_time_server', '_no_powerboard',
     )
 
     def __init__(
@@ -54,16 +60,30 @@ class Robot:
         ignored_arduinos: Optional[list[str]] = None,
         manual_boards: Optional[dict[str, list[str]]] = None,
         raw_ports: Optional[list[tuple[str, int]]] = None,
+        no_powerboard: bool = False,
     ) -> None:
-        self._lock = obtain_lock()
+        self._lock: TimeServer | socket | None
+        if IN_SIMULATOR:
+            self._lock = TimeServer.initialise()
+            if self._lock is None:
+                raise OSError(
+                    'Unable to obtain lock, Is another robot instance already running?'
+                )
+        else:
+            self._lock = obtain_lock()
         self._metadata: Optional[Metadata] = None
+        self._no_powerboard = no_powerboard
 
         setup_logging(debug, trace_logging)
         ensure_atexit_on_term()
 
         logger.info(f"sr.robot3 version {__version__}")
 
-        self._mqtt, self._astoria = init_astoria_mqtt()
+        if IN_SIMULATOR:
+            self._mqtt = None
+            self._astoria = None
+        else:
+            self._mqtt, self._astoria = init_astoria_mqtt()
 
         if manual_boards:
             self._init_power_board(manual_boards.get(PowerBoard.get_board_type(), []))
@@ -87,7 +107,10 @@ class Robot:
             defaults to None
         :raises RuntimeError: If exactly one PowerBoard is not found
         """
-        power_boards = PowerBoard._get_supported_boards(manual_boards)
+        if self._no_powerboard:
+            return
+
+        power_boards = PowerBoard._get_supported_boards(manual_boards, sleep_fn=self.sleep)
         self._power_board = singular(power_boards)
 
         # Enable all the outputs, so that we can find other boards.
@@ -108,6 +131,8 @@ class Robot:
 
         :param manual_boards:  A dictionary of board types to a list of additional
             serial port paths that should be checked for boards of that type, defaults to None
+        :param ignored_arduinos: A list of Arduino serial numbers to avoid connecting to
+        :param raw_ports: A list of serial number, baudrate tuples to try connecting to
         """
         self._raw_ports: MappingProxyType[str, RawSerial] = MappingProxyType({})
         if manual_boards is None:
@@ -122,7 +147,10 @@ class Robot:
         self._servo_boards = ServoBoard._get_supported_boards(manual_servoboards)
         self._arduinos = Arduino._get_supported_boards(manual_arduinos, ignored_arduinos)
         if raw_ports:
-            self._raw_ports = RawSerial._get_supported_boards(raw_ports)
+            if not IN_SIMULATOR:
+                self._raw_ports = RawSerial._get_supported_boards(raw_ports)
+            else:
+                logger.warning("Raw ports are not available in the simulator.")
 
     def _init_camera(self) -> None:
         """
@@ -131,9 +159,15 @@ class Robot:
         These cameras are used for AprilTag detection and provide location data of
         markers in its field of view.
         """
+        if IN_SIMULATOR:
+            publish_fn = None
+        else:
+            assert self._mqtt is not None
+            publish_fn = self._mqtt.wrapped_publish
+
         self._cameras = MappingProxyType(_setup_cameras(
             game_specific.MARKER_SIZES,
-            self._mqtt.wrapped_publish,
+            publish_fn,
         ))
 
     def _log_connected_boards(self) -> None:
@@ -142,8 +176,10 @@ class Robot:
 
         Firmware versions are also logged at debug level.
         """
+        # we only have one power board so make it iterable
+        power_board = [] if self._no_powerboard else [self.power_board]
         boards = itertools.chain(
-            [self.power_board],  # we only have one power board so make it iterable
+            power_board,
             self.motor_boards.values(),
             self.servo_boards.values(),
             self.arduinos.values(),
@@ -175,6 +211,9 @@ class Robot:
 
         :return: The power board object
         """
+        if self._no_powerboard:
+            raise RuntimeError("No power board was initialized")
+
         return self._power_board
 
     @property
@@ -276,27 +315,35 @@ class Robot:
         """
         Sleep for a number of seconds.
 
-        This is a convenience method that can be used instead of time.sleep().
-        This exists for compatibility with the simulator API only.
+        This is a convenience method that should be used instead of time.sleep()
+        to make your code compatible with the simulator.
 
         :param secs: The number of seconds to sleep for
         """
-        time.sleep(secs)
+        if IN_SIMULATOR:
+            assert isinstance(self._lock, TimeServer)
+            self._lock.sleep(secs)
+        else:
+            time.sleep(secs)
 
     @log_to_debug
     def time(self) -> float:
         """
         Get the number of seconds since the Unix Epoch.
 
-        This is a convenience method that can be used instead of time.time().
-        This exists for compatibility with the simulator API only.
+        This is a convenience method that should be used instead of time.time()
+        to make your code compatible with the simulator.
 
         NOTE: The robot's clock resets each time the robot is restarted, so this
         will not be the correct actual time but can be used to measure elapsed time.
 
         :returns: the number of seconds since the Unix Epoch.
         """
-        return time.time()
+        if IN_SIMULATOR:
+            assert isinstance(self._lock, TimeServer)
+            return self._lock.get_time()
+        else:
+            return time.time()
 
     @property
     @log_to_debug
@@ -346,9 +393,12 @@ class Robot:
         The path of the USB code drive.
 
         :returns: path to the mountpoint of the USB code drive.
-        :raises MetadataNotReadyError: If the start button has not been pressed yet
         """
-        return self._astoria.fetch_mount_path()
+        if IN_SIMULATOR:
+            return Path(os.environ['SBOT_USBKEY_PATH'])
+        else:
+            assert self._astoria is not None
+            return self._astoria.fetch_mount_path()
 
     @property
     def is_simulated(self) -> bool:
@@ -357,7 +407,7 @@ class Robot:
 
         :returns: True if the robot is simulated. False otherwise.
         """
-        return False
+        return IN_SIMULATOR
 
     @log_to_debug
     def wait_start(self) -> None:
@@ -369,25 +419,47 @@ class Robot:
         Once the start button is pressed, the metadata will be loaded and the timeout
         will start if in competition mode.
         """
+        def null_button_pressed() -> bool:
+            return False
+
+        if self._no_powerboard:
+            # null out the start button function
+            start_button_pressed = null_button_pressed
+        else:
+            start_button_pressed = self.power_board._start_button
+
+        if IN_SIMULATOR:
+            remote_start_pressed = null_button_pressed
+        else:
+            assert self._astoria is not None
+            remote_start_pressed = self._astoria.get_start_button_pressed
+
         # ignore previous button presses
-        _ = self.power_board._start_button()
-        _ = self._astoria.get_start_button_pressed()
+        _ = start_button_pressed()
+        _ = remote_start_pressed()
         logger.info('Waiting for start button.')
 
-        self.power_board.piezo.buzz(Note.A6, 0.1)
-        self.power_board._run_led.flash()
+        if not self._no_powerboard:
+            self.power_board.piezo.buzz(Note.A6, 0.1)
+            self.power_board._run_led.flash()
         self.kch._flash_start()
 
-        while not (
-            self.power_board._start_button() or self._astoria.get_start_button_pressed()
-        ):
-            time.sleep(0.1)
+        while not start_button_pressed() or remote_start_pressed():
+            self.sleep(0.1)
         logger.info("Start signal received; continuing.")
-        self.power_board._run_led.on()
+        if not self._no_powerboard:
+            self.power_board._run_led.on()
         self.kch._start = False
 
-        # Load the latest metadata that we have received over MQTT
-        self._metadata = self._astoria.fetch_current_metadata()
+        if IN_SIMULATOR:
+            self._metadata = Metadata.parse_file(
+                Path(os.environ['SBOT_METADATA_PATH']) / 'astoria.json'
+            )
+        else:
+            assert self._astoria is not None
+            # Load the latest metadata that we have received over MQTT
+            self._metadata = self._astoria.fetch_current_metadata()
 
-        if self._metadata.game_timeout is not None:
+        # Simulator timeout is handled by the simulator supervisor
+        if self._metadata.game_timeout and not IN_SIMULATOR:
             timeout.kill_after_delay(self._metadata.game_timeout)

sr/robot3/serial_wrapper.py

@@ -17,7 +17,7 @@ import serial
 
 from .exceptions import BoardDisconnectionError
 from .logging import TRACE
-from .utils import BoardIdentity
+from .utils import IN_SIMULATOR, BoardIdentity
 
 logger = logging.getLogger(__name__)
 
@@ -29,6 +29,12 @@ Param = ParamSpec("Param")
 RetType = TypeVar("RetType")
 
 E = TypeVar("E", bound=BaseException)
+BASE_TIMEOUT: float | None
+
+if IN_SIMULATOR:
+    BASE_TIMEOUT = None  # Disable timeouts while in the simulator to allow for pausing
+else:
+    BASE_TIMEOUT = 0.5
 
 
 def retry(
@@ -41,7 +47,7 @@ def retry(
     This outer function is used to pass arguments to the decorator.
 
     :param times: The number of times to retry the function.
-    :param exceptions: The exception to catch and retry on.
+    :param exceptions: The exceptions to catch and retry on.
     :return: The templated decorator function.
     """
     def decorator(func: Callable[Param, RetType]) -> Callable[Param, RetType]:
@@ -80,7 +86,7 @@ class SerialWrapper:
         self,
         port: str,
         baud: int,
-        timeout: float = 0.5,
+        timeout: float | None = BASE_TIMEOUT,
         identity: BoardIdentity = BoardIdentity(),
         delay_after_connect: float = 0,
     ):
@@ -96,6 +102,7 @@ class SerialWrapper:
             port,
             baudrate=baud,
             timeout=timeout,
+            write_timeout=timeout,
             do_not_open=True,
         )
 
@@ -124,19 +131,21 @@ class SerialWrapper:
         up to 3 times on serial errors.
 
         :param data: The data to write to the board.
+        :param endl: The endline character to add to outgoing data and remove from
+            received data
         :raises BoardDisconnectionError: If the serial connection fails during the transaction,
             including failing to respond to the command.
-        :return: The response from the board with the trailing newline removed.
+        :return: The response from the board with the trailing endline character removed.
         """
         with self._lock:
             if not self.serial.is_open:
                 if not self._connect():
                     # If the serial port cannot be opened raise an error,
                     # this will be caught by the retry decorator
-                    raise BoardDisconnectionError((
+                    raise BoardDisconnectionError(
                         f'Connection to board {self.identity.board_type}:'
                         f'{self.identity.asset_tag} could not be established',
-                    ))
+                    )
 
             try:
                 if data is not None:
@@ -172,10 +181,10 @@ class SerialWrapper:
 
             if response_str.startswith('NACK'):
                 _, error_msg = response_str.split(':', maxsplit=1)
-                logger.error((
+                logger.error(
                     f'Board {self.identity.board_type}:{self.identity.asset_tag} '
                     f'returned NACK on write command: {error_msg}'
-                ))
+                )
                 raise RuntimeError(error_msg)
 
             return response_str
@@ -185,6 +194,8 @@ class SerialWrapper:
         Send a command to the board that does not require a response.
 
         :param data: The data to write to the board.
+        :param endl: The endline character to add to outgoing data and remove from
+            received data
         :raises RuntimeError: If the board returns a NACK response,
             the firmware's error message is raised.
         """
@@ -202,15 +213,16 @@ class SerialWrapper:
         """
         try:
             self.serial.open()
-            # Wait for the board to be ready to receive data
-            # Certain boards will reset when the serial port is opened
-            time.sleep(self.delay_after_connect)
+            if not IN_SIMULATOR:
+                # Wait for the board to be ready to receive data
+                # Certain boards will reset when the serial port is opened
+                time.sleep(self.delay_after_connect)
             self.serial.reset_input_buffer()
         except serial.SerialException:
-            logger.error((
+            logger.error(
                 'Failed to connect to board '
                 f'{self.identity.board_type}:{self.identity.asset_tag}'
-            ))
+            )
             return False
 
         logger.info(

sr/robot3/servo_board.py

@@ -12,14 +12,15 @@ 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,
 )
 
 DUTY_MIN = 500
 DUTY_MAX = 4000
 START_DUTY_MIN = 1000
 START_DUTY_MAX = 2000
+NUM_SERVOS = 12
 
 logger = logging.getLogger(__name__)
 BAUDRATE = 115200  # Since the servo board is a USB device, this is ignored
@@ -77,7 +78,7 @@ class ServoBoard(Board):
         self._serial = SerialWrapper(serial_port, BAUDRATE, identity=initial_identity)
 
         self._servos = tuple(
-            Servo(self._serial, index) for index in range(12)
+            Servo(self._serial, index) for index in range(NUM_SERVOS)
         )
 
         self._identity = self.identify()
@@ -109,17 +110,46 @@ class ServoBoard(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 servo board at port {serial_port!r}, "
-                    "could not be identified. Ignoring this device")
-            else:
-                logger.warning(
-                    f"Found servo 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 servo board at port {serial_port!r}, "
+                        "could not be identified. Ignoring this device")
+                elif initial_identity.manufacturer == 'sbot_simulator':
+                    logger.warning(
+                        f"Simulator specified servo board at port {serial_port!r}, "
+                        "could not be identified. Ignoring this device")
+                return None
+
+            logger.warning(
+                f"Found servo 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, ServoBoard]:
+        """
+        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('ServoBoard'):
+
+            # 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,
@@ -133,6 +163,9 @@ class ServoBoard(Board):
             to connect to, defaults to None
         :return: A mapping of serial numbers to servo boards.
         """
+        if IN_SIMULATOR:
+            return cls._get_simulator_boards()
+
         boards = {}
         serial_ports = comports()
         for port in serial_ports: