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

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__)
 
@@ -40,10 +43,12 @@ class Robot:
     :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__(
@@ -55,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(), []))
@@ -88,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.
@@ -125,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:
         """
@@ -134,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:
@@ -145,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(),
@@ -178,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
@@ -279,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
@@ -350,7 +394,11 @@ class Robot:
 
         :returns: path to the mountpoint of the USB code drive.
         """
-        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:
@@ -359,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:
@@ -371,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(
@@ -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,
         )
 
@@ -206,9 +213,10 @@ 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(

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:

sr/robot3/simulator/camera.py

@@ -0,0 +1,75 @@
+import struct
+
+import cv2
+import numpy as np
+from april_vision import FrameSource
+from numpy.typing import NDArray
+from serial import serial_for_url
+
+from ..utils import BoardInfo
+
+HEADER_SIZE = 5  # 1 byte for the type, 4 bytes for the length
+IMAGE_TAG_ID = 0
+
+
+class WebotsRemoteCameraSource(FrameSource):
+    # Webots cameras include an alpha channel, this informs april_vision of how to handle it
+    COLOURSPACE = cv2.COLOR_BGRA2GRAY
+
+    def __init__(self, camera_info: BoardInfo) -> None:
+        self.calibration = (0.0, 0.0, 0.0, 0.0)
+        # Use pyserial to give a nicer interface for connecting to the camera socket
+        self._serial = serial_for_url(camera_info.url, baudrate=115200, timeout=None)
+
+        # Check the camera is connected
+        response = self._make_request("*IDN?")
+        if not response.split(b":")[1].lower().startswith(b"cam"):
+            raise RuntimeError(f"Camera not connected to a camera, returned: {response!r}")
+
+        # Get the calibration data for this camera
+        response = self._make_request("CAM:CALIBRATION?")
+
+        # The calibration data is returned as a string of floats separated by colons
+        new_calibration = tuple(map(float, response.split(b":")))
+        assert len(new_calibration) == 4, f"Invalid calibration data: {new_calibration}"
+        self.calibration = new_calibration
+
+        # Get the image size for this camera
+        response = self._make_request("CAM:RESOLUTION?")
+        self.image_size = tuple(map(int, response.split(b":")))
+        assert len(self.image_size) == 2, f"Invalid image dimensions: {self.image_size}"
+
+    def read(self, fresh: bool = True) -> NDArray:
+        """
+        The method for getting a new frame.
+
+        :param fresh: Whether to flush the device's buffer before capturing
+        the frame, unused.
+        """
+        self._serial.write(b"CAM:FRAME!\n")
+        # The image is encoded as a TLV (Type, Length, Value) packet
+        # so we need to read the header to get the type and length of the image
+        header = self._serial.read(HEADER_SIZE)
+        assert len(header) == HEADER_SIZE, f"Invalid header length: {len(header)}"
+        img_tag, img_len = struct.unpack('>BI', header)
+        assert img_tag == IMAGE_TAG_ID, f"Invalid image tag: {img_tag}"
+
+        # Get the image data now we know the length
+        img_data = self._serial.read(img_len)
+        assert len(img_data) == img_len, f"Invalid image data length: {len(img_data)}"
+
+        rgb_frame_raw: NDArray[np.uint8] = np.frombuffer(img_data, np.uint8)
+
+        # Height is first, then width, then channels
+        return rgb_frame_raw.reshape((self.image_size[1], self.image_size[0], 4))
+
+    def close(self) -> None:
+        """Close the underlying socket on exit."""
+        self._serial.close()
+
+    def _make_request(self, command: str) -> bytes:
+        self._serial.write(command.encode() + b"\n")
+        response = self._serial.readline()
+        if not response.endswith(b"\n") or response.startswith(b"NACK:"):
+            raise RuntimeError(f"Failed to communicate with camera, returned: {response!r}")
+        return response

sr/robot3/simulator/time_server.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+
+from ..exceptions import BoardDisconnectionError, IncorrectBoardError
+from ..serial_wrapper import SerialWrapper
+from ..utils import BoardIdentity, get_simulator_boards
+
+logger = logging.getLogger(__name__)
+
+BAUDRATE = 115200
+
+
+class TimeServer:
+    @staticmethod
+    def get_board_type() -> str:
+        """
+        Return the type of the board.
+
+        :return: The literal string 'TimeServer'.
+        """
+        return 'TimeServer'
+
+    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,
+            # Disable the timeout so sleep works properly
+            timeout=None,
+        )
+
+        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)
+
+    @classmethod
+    def initialise(cls) -> 'TimeServer' | None:
+        # The filter here is the name of the emulated board in the simulator
+        boards = get_simulator_boards('TimeServer')
+
+        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 time server 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 get_time(self) -> float:
+        time_str = self._serial.query('TIME?')
+        return datetime.fromisoformat(time_str).timestamp()
+
+    def sleep(self, duration: float) -> None:
+        if duration < 0:
+            raise ValueError("sleep length must be non-negative")
+
+        self._serial.query(f'SLEEP:{int(duration * 1000)}')

sr/robot3/timeout.py

@@ -1,4 +1,5 @@
 """Functions for killing the robot after a certain amount of time."""
+import atexit
 import logging
 import os
 import signal
@@ -9,6 +10,9 @@ from typing import Optional
 
 logger = logging.getLogger(__name__)
 
+TIMEOUT_MESSAGE = "Timeout expired: Game Over!"
+EXIT_ATTEMPTS = 0
+
 
 def timeout_handler(signal_type: int, stack_frame: Optional[FrameType]) -> None:
     """
@@ -17,13 +21,34 @@ def timeout_handler(signal_type: int, stack_frame: Optional[FrameType]) -> None:
     This function is called when the timeout expires and will stop the robot's main process.
     In order for this to work, any threads that are created must be daemons.
 
+    If the process doesn't terminate clearly (perhaps because the exception was caught),
+    exit less cleanly.
+
     NOTE: This function is not called on Windows.
 
     :param signal_type: The signal that triggered this handler
     :param stack_frame: The stack frame at the time of the signal
     """
-    logger.info("Timeout expired: Game Over!")
-    exit(0)
+    global EXIT_ATTEMPTS
+    EXIT_ATTEMPTS += 1
+
+    if sys.platform == "win32":
+        raise AssertionError("This function should not be called on Windows")
+
+    if EXIT_ATTEMPTS == 1:
+        # Allow 2 seconds for the process to exit cleanly before killing it
+        signal.alarm(2)
+        logger.info(TIMEOUT_MESSAGE)
+        # Exit cleanly
+        exit(0)
+    else:
+        # The process didn't exit cleanly, so first call the cleanup handlers
+        # and use an unhanded alarm to force python to exit with a core dump
+        signal.signal(signal.SIGALRM, signal.SIG_DFL)
+        signal.alarm(2)  # Allow 2 seconds for cleanup
+
+        atexit._run_exitfuncs()
+        exit(0)
 
 
 def win_timeout_handler() -> None:
@@ -35,7 +60,7 @@ def win_timeout_handler() -> None:
 
     NOTE: This function is only called on Windows.
     """
-    logger.info("Timeout expired: Game Over!")
+    logger.info(TIMEOUT_MESSAGE)
     os.kill(os.getpid(), signal.SIGTERM)
 
 
@@ -52,6 +77,7 @@ def kill_after_delay(timeout_seconds: int) -> None:
         # Windows doesn't have SIGALRM,
         # so we approximate its functionality using a delayed SIGTERM
         timer = Timer(timeout_seconds, win_timeout_handler)
+        timer.daemon = True
         timer.start()
     else:
         signal.signal(signal.SIGALRM, timeout_handler)