From c95ba66c1ac64f93c593c1fb8573afb228a020ab Mon Sep 17 00:00:00 2001 From: Mike Silva Date: Mon, 19 Oct 2026 20:54:31 -0400 Subject: [PATCH 1/2] Initial commit of tankdrive.py and test --- .../templates/components/tankdrive.py | 286 ++++++++++++++++++ tests/test_tankdrive.py | 252 +++++++++++++++ 2 files changed, 538 insertions(+) create mode 100644 src/pybot_sdk/templates/components/tankdrive.py create mode 100644 tests/test_tankdrive.py diff --git a/src/pybot_sdk/templates/components/tankdrive.py b/src/pybot_sdk/templates/components/tankdrive.py new file mode 100644 index 0000000..c7f618a --- /dev/null +++ b/src/pybot_sdk/templates/components/tankdrive.py @@ -0,0 +1,286 @@ +"""A tank-drive component for MagicBot. + +Put the motors and any sensors you want to use in a dictionary. For example: + + import wpilib + from components import XboxController + from components.gyro import Gyro + from components.tankdrive import Tankdrive + from magicbot import MagicRobot + + class Robot(MagicRobot): + gyro: Gyro + drivetrain: Tankdrive + controller: XboxController + + def createObjects(self) -> None: + left_motors = wpilib.MotorControllerGroup( + wpilib.PWMSparkMax(0), wpilib.PWMSparkMax(1) + ) + right_motors = wpilib.MotorControllerGroup( + wpilib.PWMSparkMax(2), wpilib.PWMSparkMax(3) + ) + self.drivetrain_components = { + "left_motor": left_motors, + "right_motor": right_motors, + "left_encoder": wpilib.Encoder(0, 1), + "right_encoder": wpilib.Encoder(2, 3), + "gyro": self.gyro, + } + self.drivetrain_track_width = 0.62 + self.drivetrain_units = "meters" + self.drivetrain_drive_mode = "arcade" + + def teleopPeriodic(self) -> None: + self.drivetrain.throttle = -self.controller.getLeftY() + self.drivetrain.turn = self.controller.getRightX() + +To read the robot's position in feet, call +``self.drivetrain.get_pose("feet")``. The default is meters. + +The dictionary must include ``left_motor`` and ``right_motor``. You can use a +single motor or a group of motors on each side. Encoders and the gyro are +optional. The robot's position is tracked when both encoders and a gyro are +provided. Otherwise, ``pose`` stays at its starting position. The gyro should +be the ``Gyro`` component from ``gyro.py``. + +Set ``track_width`` to the distance between the left and right wheels. +``units`` defaults to ``"meters"`` and tells Tankdrive which units you use for +the track width and encoder readings. Supported units are meters, inches, and +feet. The ``get_pose`` method can return the robot's position in any of these +units. Tankdrive stores its pose in meters, as WPILib expects. + +Set ``drive_mode`` to ``"arcade"``, ``"cheesy"``, or ``"tank"``. Arcade and +Cheesy use ``throttle`` and ``turn``. Tank mode uses ``left`` and ``right``. +Cheesy mode also uses the ``quick_turn`` input. Set these inputs from your +robot code before Tankdrive runs. + +In simulation, call +``set_simulation_readings(left_distance, right_distance, heading)`` if your +physics model supplies sensor values directly. Give distances in the units +you selected above. +""" + +from __future__ import annotations + +from typing import Any + +import wpilib +import wpilib.drive +from magicbot import feedback +from wpimath import units as wpimath_units +from wpimath.geometry import Pose2d, Rotation2d +from wpimath.kinematics import DifferentialDriveKinematics, DifferentialDriveOdometry + + +class Tankdrive: + """Drive a robot using arcade, Cheesy, or tank controls.""" + + # MagicBot injects these from robot attributes named after the component. + # For a component named "drivetrain", set drivetrain_components, etc. + components: dict[str, Any] + track_width: float + units: str = "meters" + drive_mode: str = "arcade" + + throttle: float = 0.0 + turn: float = 0.0 + left: float = 0.0 + right: float = 0.0 + quick_turn: bool = False + + def setup(self) -> None: + """Set up motor control and position tracking.""" + if not isinstance(self.components, dict): + raise ValueError("tankdrive_components must be a dictionary") + + try: + left_motor = self.components["left_motor"] + right_motor = self.components["right_motor"] + except KeyError as error: + raise ValueError( + "tankdrive_components must contain 'left_motor' and 'right_motor'" + ) from error + + self.drive = wpilib.drive.DifferentialDrive(left_motor, right_motor) + self.is_simulation = wpilib.RobotBase.isSimulation() + track_width = self._convert_to_internal_units(self.track_width) + if track_width <= 0: + raise ValueError("tankdrive_track_width must be greater than zero") + self.kinematics = DifferentialDriveKinematics(track_width) + self.pose = Pose2d() + self.odometry: DifferentialDriveOdometry | None = None + + self._left_encoder = self.components.get("left_encoder") + self._right_encoder = self.components.get("right_encoder") + self._gyro = self.components.get("gyro") + self._simulation_readings: tuple[float, float, Rotation2d] | None = None + + if ( + self._left_encoder is not None + and self._right_encoder is not None + and self._gyro is not None + ): + self.odometry = DifferentialDriveOdometry( + self._get_heading(), + self._get_left_distance(), + self._get_right_distance(), + ) + + @staticmethod + def _convert_distance(distance: float, from_units: str, to_units: str) -> float: + """Convert a distance between meters, inches, or feet.""" + from_units = from_units.strip().lower() + to_units = to_units.strip().lower() + + if from_units in ("m", "meter", "meters"): + distance_in_meters = distance + elif from_units in ("in", "inch", "inches"): + distance_in_meters = wpimath_units.inchesToMeters(distance) + elif from_units in ("ft", "foot", "feet"): + distance_in_meters = wpimath_units.feetToMeters(distance) + else: + raise ValueError( + f"Unsupported distance unit {from_units!r}; use meters, inches, or feet" + ) + + if to_units in ("m", "meter", "meters"): + return distance_in_meters + if to_units in ("in", "inch", "inches"): + return wpimath_units.metersToInches(distance_in_meters) + if to_units in ("ft", "foot", "feet"): + return wpimath_units.metersToFeet(distance_in_meters) + raise ValueError( + f"Unsupported distance unit {to_units!r}; use meters, inches, or feet" + ) + + def _convert_to_internal_units(self, distance: float) -> float: + """Convert one of the configured distances into WPILib's standard unit.""" + return self._convert_distance(distance, self.units, "meters") + + def _get_heading(self) -> Rotation2d: + """Get the robot's heading from the gyro or simulation.""" + if self._simulation_readings is not None: + return self._simulation_readings[2] + if self._gyro is None: + raise RuntimeError("Odometry requires a gyro") + return self._gyro.rotation + + def _get_left_distance(self) -> float: + """Get the left encoder distance in meters.""" + if self._simulation_readings is not None: + return self._convert_to_internal_units(self._simulation_readings[0]) + if self._left_encoder is None: + raise RuntimeError("Odometry requires a left encoder") + return self._convert_to_internal_units(self._left_encoder.getDistance()) + + def _get_right_distance(self) -> float: + """Get the right encoder distance in meters.""" + if self._simulation_readings is not None: + return self._convert_to_internal_units(self._simulation_readings[1]) + if self._right_encoder is None: + raise RuntimeError("Odometry requires a right encoder") + return self._convert_to_internal_units(self._right_encoder.getDistance()) + + def execute(self) -> None: + """Drive the robot and update its position if sensors are available.""" + mode = self.drive_mode.strip().lower() + if mode == "arcade": + self.drive.arcadeDrive(self.throttle, self.turn) + elif mode == "cheesy": + self.drive.curvatureDrive(self.throttle, self.turn, self.quick_turn) + elif mode == "tank": + self.drive.tankDrive(self.left, self.right) + else: + raise ValueError( + f"Unknown tankdrive_drive_mode {self.drive_mode!r}; " + "valid options: 'arcade', 'cheesy', or 'tank'" + ) + + if self.odometry is not None: + self.pose = self.odometry.update( + self._get_heading(), + self._get_left_distance(), + self._get_right_distance(), + ) + + def reset_odometry(self, pose: Pose2d = Pose2d()) -> None: + """Reset the estimated pose to ``pose`` using the current sensors. + + Raises: + RuntimeError: If encoders and a gyro were not all configured. + + """ + if self.odometry is None: + raise RuntimeError("Odometry requires two encoders and a gyro") + self.odometry.resetPosition( + self._get_heading(), + self._get_left_distance(), + self._get_right_distance(), + pose, + ) + self.pose = pose + + def set_simulation_readings( + self, + left_distance: float, + right_distance: float, + heading: Rotation2d, + ) -> None: + """Set simulated sensor distances and heading for odometry. + + Distances use the configured ``units``. This method does not command + the motors; a physics model remains responsible for motor simulation. + + Raises: + RuntimeError: If called on a real robot or without odometry. + + """ + if not self.is_simulation: + raise RuntimeError("Simulation readings require a simulated robot") + if self.odometry is None: + raise RuntimeError("Simulation readings require odometry sensors") + self._simulation_readings = (left_distance, right_distance, heading) + + def reset_encoders(self) -> None: + """Reset the configured encoders, if both are available.""" + if self._left_encoder is not None: + self._left_encoder.reset() + if self._right_encoder is not None: + self._right_encoder.reset() + + def stop(self) -> None: + """Stop both sides of the drivetrain.""" + self.drive.stopMotor() + + def get_pose(self, units: str = "meters") -> Pose2d: + """Return the current robot position in the requested distance units. + + The stored ``pose`` always uses meters. This method returns a new pose + with its X and Y coordinates converted to meters, inches, or feet. + """ + return Pose2d( + self._convert_distance(self.pose.X(), "meters", units), + self._convert_distance(self.pose.Y(), "meters", units), + self.pose.rotation(), + ) + + @feedback + def odometry_enabled(self) -> bool: + """Show whether pose estimation has all required sensors.""" + return self.odometry is not None + + @feedback + def pose_x(self) -> float: + """Show the estimated robot X coordinate in meters.""" + return self.pose.X() + + @feedback + def pose_y(self) -> float: + """Show the estimated robot Y coordinate in meters.""" + return self.pose.Y() + + @feedback + def heading_degrees(self) -> float: + """Show the estimated robot heading in degrees.""" + return self.pose.rotation().degrees() diff --git a/tests/test_tankdrive.py b/tests/test_tankdrive.py new file mode 100644 index 0000000..98e757e --- /dev/null +++ b/tests/test_tankdrive.py @@ -0,0 +1,252 @@ +"""Tests for the packaged Tankdrive MagicBot component.""" + +import sys +from importlib.util import module_from_spec, spec_from_file_location +from math import isclose +from pathlib import Path +from types import ModuleType, SimpleNamespace + + +def _load_tankdrive_module(monkeypatch, simulation: bool = False): + """Load the template against focused WPILib fakes.""" + calls = [] + + class FakeRotation2d: + def __init__(self, degrees: float = 0.0): + self._degrees = degrees + + def degrees(self) -> float: + return self._degrees + + class FakePose2d: + def __init__(self, x: float = 0.0, y: float = 0.0, rotation=None): + self._x = x + self._y = y + self._rotation = rotation or FakeRotation2d() + + def X(self) -> float: + return self._x + + def Y(self) -> float: + return self._y + + def rotation(self): + return self._rotation + + class FakeDifferentialDrive: + def __init__(self, left_motor, right_motor): + calls.append(("init_drive", left_motor, right_motor)) + + def arcadeDrive(self, throttle, turn): + calls.append(("arcade", throttle, turn)) + + def curvatureDrive(self, throttle, turn, quick_turn): + calls.append(("curvature", throttle, turn, quick_turn)) + + def tankDrive(self, left, right): + calls.append(("tank", left, right)) + + def stopMotor(self): + calls.append(("stop",)) + + class FakeDifferentialDriveKinematics: + def __init__(self, track_width): + calls.append(("kinematics", track_width)) + + class FakeDifferentialDriveOdometry: + def __init__(self, heading, left_distance, right_distance): + calls.append(("init_odometry", heading, left_distance, right_distance)) + + def update(self, heading, left_distance, right_distance): + calls.append(("update_odometry", heading, left_distance, right_distance)) + return FakePose2d(left_distance, right_distance, heading) + + def resetPosition(self, heading, left_distance, right_distance, pose): + calls.append( + ("reset_odometry", heading, left_distance, right_distance, pose) + ) + + wpilib = ModuleType("wpilib") + wpilib.__path__ = [] + wpilib.RobotBase = SimpleNamespace(isSimulation=lambda: simulation) + drive_module = ModuleType("wpilib.drive") + drive_module.DifferentialDrive = FakeDifferentialDrive + wpilib.drive = drive_module + + magicbot = ModuleType("magicbot") + magicbot.feedback = lambda function: function + + wpimath = ModuleType("wpimath") + wpimath.__path__ = [] + geometry = ModuleType("wpimath.geometry") + geometry.Pose2d = FakePose2d + geometry.Rotation2d = FakeRotation2d + kinematics = ModuleType("wpimath.kinematics") + kinematics.DifferentialDriveKinematics = FakeDifferentialDriveKinematics + kinematics.DifferentialDriveOdometry = FakeDifferentialDriveOdometry + units = ModuleType("wpimath.units") + units.inchesToMeters = lambda inches: inches * 0.0254 + units.feetToMeters = lambda feet: feet * 0.3048 + units.metersToInches = lambda meters: meters / 0.0254 + units.metersToFeet = lambda meters: meters / 0.3048 + + monkeypatch.setitem(sys.modules, "wpilib", wpilib) + monkeypatch.setitem(sys.modules, "wpilib.drive", drive_module) + monkeypatch.setitem(sys.modules, "magicbot", magicbot) + monkeypatch.setitem(sys.modules, "wpimath", wpimath) + monkeypatch.setitem(sys.modules, "wpimath.geometry", geometry) + monkeypatch.setitem(sys.modules, "wpimath.kinematics", kinematics) + monkeypatch.setitem(sys.modules, "wpimath.units", units) + wpimath.units = units + + template_path = ( + Path(__file__).resolve().parents[1] + / "src" + / "pybot_sdk" + / "templates" + / "components" + / "tankdrive.py" + ) + spec = spec_from_file_location("tankdrive_template", template_path) + assert spec is not None + assert spec.loader is not None + module = module_from_spec(spec) + spec.loader.exec_module(module) + return module, calls, FakePose2d, FakeRotation2d + + +class FakeEncoder: + """Minimal encoder fake with a configurable distance.""" + + def __init__(self, distance: float): + """Store the current distance.""" + self.distance = distance + self.reset_count = 0 + + def getDistance(self) -> float: + """Return the configured distance.""" + return self.distance + + def reset(self) -> None: + """Count resets for test assertions.""" + self.reset_count += 1 + + +class FakeGyro: + """Minimal gyro fake with a fixed rotation.""" + + def __init__(self, rotation): + """Store the current rotation.""" + self.rotation = rotation + + +def _make_drivetrain(module, track_width=24.0, units="in", sensors=True): + left_motor = object() + right_motor = object() + components = { + "left_motor": left_motor, + "right_motor": right_motor, + } + if sensors: + components.update( + left_encoder=FakeEncoder(12.0), + right_encoder=FakeEncoder(24.0), + gyro=FakeGyro(module.Rotation2d(30.0)), + ) + drivetrain = module.Tankdrive() + drivetrain.components = components + drivetrain.track_width = track_width + drivetrain.units = units + return drivetrain, left_motor, right_motor + + +def test_setup_builds_kinematics_odometry_and_drives_each_mode(monkeypatch) -> None: + """Use configured units for geometry and dispatch all control modes.""" + module, calls, _, _ = _load_tankdrive_module(monkeypatch) + drivetrain, left_motor, right_motor = _make_drivetrain(module) + drivetrain.setup() + + assert calls[0] == ("init_drive", left_motor, right_motor) + assert calls[1][0] == "kinematics" + assert isclose(calls[1][1], 0.6096) + assert calls[2][0] == "init_odometry" + assert calls[2][1] is drivetrain._gyro.rotation + assert isclose(calls[2][2], 0.3048) + assert isclose(calls[2][3], 0.6096) + + drivetrain.throttle = 0.5 + drivetrain.turn = -0.25 + drivetrain.quick_turn = True + drivetrain.execute() + drivetrain.drive_mode = "cheesy" + drivetrain.execute() + drivetrain.drive_mode = "tank" + drivetrain.left = 0.2 + drivetrain.right = 0.4 + drivetrain.execute() + + assert [call for call in calls if call[0] in ("arcade", "curvature", "tank")] == [ + ("arcade", 0.5, -0.25), + ("curvature", 0.5, -0.25, True), + ("tank", 0.2, 0.4), + ] + assert isclose(drivetrain.pose.X(), 0.3048) + assert isclose(drivetrain.pose.Y(), 0.6096) + pose_in_inches = drivetrain.get_pose("inches") + assert isclose(pose_in_inches.X(), 12.0) + assert isclose(pose_in_inches.Y(), 24.0) + assert isclose(drivetrain.get_pose("feet").X(), 1.0) + + +def test_missing_odometry_sensors_do_not_prevent_driving(monkeypatch) -> None: + """Allow a motors-only configuration while keeping a pose attribute.""" + module, calls, _, _ = _load_tankdrive_module(monkeypatch) + drivetrain, _, _ = _make_drivetrain(module, sensors=False) + drivetrain.setup() + drivetrain.execute() + + assert drivetrain.odometry is None + assert drivetrain.pose.X() == 0.0 + assert ("arcade", 0.0, 0.0) in calls + + +def test_simulation_readings_update_odometry_in_configured_units(monkeypatch) -> None: + """Allow a physics model to provide simulated sensor measurements.""" + module, calls, _, rotation_type = _load_tankdrive_module( + monkeypatch, simulation=True + ) + drivetrain, _, _ = _make_drivetrain(module) + drivetrain.setup() + heading = rotation_type(90.0) + drivetrain.set_simulation_readings(36.0, 48.0, heading) + drivetrain.execute() + + update = next(call for call in calls if call[0] == "update_odometry") + assert update[1] is heading + assert isclose(update[2], 0.9144) + assert isclose(update[3], 1.2192) + assert isclose(drivetrain.pose.X(), 0.9144) + + +def test_invalid_mode_and_missing_motors_raise_clear_errors(monkeypatch) -> None: + """Reject incomplete motor configuration and unknown drive modes.""" + module, _, _, _ = _load_tankdrive_module(monkeypatch) + drivetrain = module.Tankdrive() + drivetrain.components = {"left_motor": object()} + drivetrain.track_width = 0.6 + try: + drivetrain.setup() + except ValueError as error: + assert "left_motor" in str(error) and "right_motor" in str(error) + else: + raise AssertionError("expected missing motor configuration to fail") + + drivetrain, _, _ = _make_drivetrain(module, sensors=False) + drivetrain.setup() + drivetrain.drive_mode = "unsupported" + try: + drivetrain.execute() + except ValueError as error: + assert "unsupported" in str(error) + else: + raise AssertionError("expected an unknown drive mode to fail") From 92dc65c7d118c645f8b7c740fbb4516c0d2e00c2 Mon Sep 17 00:00:00 2001 From: Mike Silva Date: Mon, 19 Oct 2026 20:54:44 -0400 Subject: [PATCH 2/2] Updates for tankdrive --- README.md | 1 + tests/test_cli.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/README.md b/README.md index a0cbe8b..ed20336 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,7 @@ pybot new robot --template xrp # Add components to an existing MagicBot project pybot add component gyro pybot add component logger +pybot add component tankdrive ``` The MagicBot starter template targets Python 3.14 and RobotPy 2026. Its RobotPy diff --git a/tests/test_cli.py b/tests/test_cli.py index a8814c8..5fa47f8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -85,6 +85,22 @@ def test_add_logger_component_copies_template_to_project_components( assert capsys.readouterr().out == "Added component logger to components/\n" +def test_add_tankdrive_component_copies_template_to_project_components( + tmp_path: Path, monkeypatch, capsys +) -> None: + """Copy the Tankdrive component template into the project.""" + project = tmp_path / "robot_project" + components = project / "components" + components.mkdir(parents=True) + monkeypatch.chdir(project) + + assert main(["add", "component", "tankdrive"]) == 0 + assert (components / "tankdrive.py").read_bytes() == resources.files( + "pybot_sdk" + ).joinpath("templates", "components", "tankdrive.py").read_bytes() + assert capsys.readouterr().out == "Added component tankdrive to components/\n" + + def _load_logger_module( monkeypatch, data_log_manager: SimpleNamespace, entry_calls: list | None = None ):