diff --git a/examples/robot/DSGamepadChooser/robot.py b/examples/robot/DSGamepadChooser/robot.py index a33a0c74c..01bbc03ac 100644 --- a/examples/robot/DSGamepadChooser/robot.py +++ b/examples/robot/DSGamepadChooser/robot.py @@ -7,7 +7,7 @@ from hal import RobotMode import wpilib -from wpilib.opmoderobot import OpModeRobot +from wpilib.opmodes import OpModeRobot class DoNothingTeleop(wpilib.PeriodicOpMode): diff --git a/examples/robot/ExpansionHubSample/opmodes/default_auto_mode.py b/examples/robot/ExpansionHubSample/opmodes/default_auto_mode.py new file mode 100644 index 000000000..6cdae02e2 --- /dev/null +++ b/examples/robot/ExpansionHubSample/opmodes/default_auto_mode.py @@ -0,0 +1,33 @@ +# Copyright (c) FIRST and other WPILib contributors. +# Open Source Software; you can modify and/or share it under the terms of +# the WPILib BSD license file in the root directory of this project. + +from typing import TYPE_CHECKING + +import wpilib + +if TYPE_CHECKING: + from robot import Robot + + +@wpilib.autonomous +class DefaultAutoMode(wpilib.PeriodicOpMode): + def __init__(self, robot: "Robot") -> None: + super().__init__() + self.robot = robot + self.timer = wpilib.Timer() + + def start(self) -> None: + self.timer.reset() + self.timer.start() + + def periodic(self) -> None: + if self.timer.get() < 2.0: + self.robot.motor0.set_throttle(0.5) + self.robot.motor1.set_throttle(0.5) + elif self.timer.get() < 4.0: + self.robot.motor0.set_throttle(0.9) + self.robot.motor1.set_throttle(0.9) + else: + self.robot.motor0.set_throttle(0.0) + self.robot.motor1.set_throttle(0.0) diff --git a/examples/robot/ExpansionHubSample/opmodes/default_tele_mode.py b/examples/robot/ExpansionHubSample/opmodes/default_tele_mode.py new file mode 100644 index 000000000..f8dcfe4a6 --- /dev/null +++ b/examples/robot/ExpansionHubSample/opmodes/default_tele_mode.py @@ -0,0 +1,26 @@ +# Copyright (c) FIRST and other WPILib contributors. +# Open Source Software; you can modify and/or share it under the terms of +# the WPILib BSD license file in the root directory of this project. + +from typing import TYPE_CHECKING + +import wpilib + +if TYPE_CHECKING: + from robot import Robot + + +@wpilib.teleop +class DefaultTeleMode(wpilib.PeriodicOpMode): + def __init__(self, robot: "Robot") -> None: + super().__init__() + self.robot = robot + self.gamepad = wpilib.DriverStation.get_gamepad(0) + + def periodic(self) -> None: + self.robot.motor0.set_throttle(-self.gamepad.get_left_y()) + self.robot.motor1.set_throttle(-self.gamepad.get_right_y()) + self.robot.motor2.set_throttle(-self.gamepad.get_left_x()) + self.robot.motor3.set_throttle(-self.gamepad.get_right_x()) + self.robot.servo0.set_position(self.gamepad.get_left_trigger()) + self.robot.servo1.set_position(self.gamepad.get_right_trigger()) diff --git a/examples/robot/ExpansionHubSample/robot.py b/examples/robot/ExpansionHubSample/robot.py new file mode 100644 index 000000000..df33f3db5 --- /dev/null +++ b/examples/robot/ExpansionHubSample/robot.py @@ -0,0 +1,23 @@ +# Copyright (c) FIRST and other WPILib contributors. +# Open Source Software; you can modify and/or share it under the terms of +# the WPILib BSD license file in the root directory of this project. + +import wpilib + + +class Robot(wpilib.OpModeRobot): + """Demo robot for Expansion Hub motors and servos. + + The motors and servos are driven using the controllers in the teleop OpMode + and timed in the autonomous OpMode. + """ + + def __init__(self) -> None: + """Called once at the beginning of the robot program.""" + super().__init__() + self.motor0 = wpilib.ExpansionHubMotor(0, 0) + self.motor1 = wpilib.ExpansionHubMotor(0, 1) + self.motor2 = wpilib.ExpansionHubMotor(0, 2) + self.motor3 = wpilib.ExpansionHubMotor(0, 3) + self.servo0 = wpilib.ExpansionHubServo(0, 0) + self.servo1 = wpilib.ExpansionHubServo(0, 1) diff --git a/examples/robot/examples.toml b/examples/robot/examples.toml index 35494cfd9..01b0f54e7 100644 --- a/examples/robot/examples.toml +++ b/examples/robot/examples.toml @@ -13,6 +13,7 @@ base = [ "ElevatorSimulation", "ElevatorTrapezoidProfile", "Encoder", + "ExpansionHubSample", "GettingStarted", "Gyro", "DSGamepadChooser", diff --git a/subprojects/robotpy-wpilib/tests/framework/test_opmode_robot.py b/subprojects/robotpy-wpilib/tests/framework/test_opmode_robot.py index ddea237df..a2ddadd4a 100644 --- a/subprojects/robotpy-wpilib/tests/framework/test_opmode_robot.py +++ b/subprojects/robotpy-wpilib/tests/framework/test_opmode_robot.py @@ -1,13 +1,1447 @@ -import pytest +import importlib +from pathlib import Path +import sys +import textwrap import threading import tunables +from typing import get_args, get_origin, get_overloads, get_type_hints + +import pytest +import wpilib +import wpilib._impl.opmode as opmode_impl +from wpilib import OpMode from wpilib import simulation as wsim -from wpilib.opmoderobot import OpModeRobot -from wpilib import OpMode, RobotState +from wpilib._wpilib import OpModeRobotBase, RobotState +from wpilib.opmodes import OpModeRobot, autonomous, teleop, utility from hal import RobotMode from wpiutil import Color +@pytest.fixture(autouse=True) +def reset_decorated_opmodes(monkeypatch): + monkeypatch.setattr(opmode_impl, "_decorated_opmodes", []) + RobotState.clear_opmodes() + yield + for module_name in tuple(sys.modules): + if ( + module_name == "samplebot" + or module_name.startswith("samplebot.") + or module_name == "robot" + or module_name == "opmodes" + or module_name.startswith("opmodes.") + ): + sys.modules.pop(module_name) + + +def import_robot_package(monkeypatch, tmp_path, robot_source, module_contents): + package_name = "samplebot" + package_dir = tmp_path / package_name + opmodes_dir = package_dir / "opmodes" + opmodes_dir.mkdir(parents=True) + (package_dir / "__init__.py").write_text("") + (package_dir / "robot.py").write_text(textwrap.dedent(robot_source)) + (opmodes_dir / "__init__.py").write_text("") + + for relative_path, source in module_contents.items(): + module_path = package_dir / relative_path + module_path.parent.mkdir(parents=True, exist_ok=True) + module_path.write_text(textwrap.dedent(source)) + + for module_name in tuple(sys.modules): + if module_name == package_name or module_name.startswith(f"{package_name}."): + sys.modules.pop(module_name) + + monkeypatch.syspath_prepend(str(tmp_path)) + importlib.invalidate_caches() + return package_dir, importlib.import_module(f"{package_name}.robot") + + +def test_opmode_decorators_attach_metadata(): + @autonomous(group="Drive", description="Auto desc") + class AutoMode(OpMode): + pass + + @teleop + class TeleMode(OpMode): + pass + + @utility( + name="Arm Test", + text_color=Color.WHITE, + background_color=Color.BLACK, + ) + class UtilityMode(OpMode): + pass + + assert AutoMode._wpilib_opmode_metadata.mode == RobotMode.AUTONOMOUS + assert AutoMode._wpilib_opmode_metadata.name == "AutoMode" + assert AutoMode._wpilib_opmode_metadata.group == "Drive" + assert AutoMode._wpilib_opmode_metadata.description == "Auto desc" + assert TeleMode._wpilib_opmode_metadata.mode == RobotMode.TELEOPERATED + assert TeleMode._wpilib_opmode_metadata.name == "TeleMode" + assert UtilityMode._wpilib_opmode_metadata.mode == RobotMode.UTILITY + assert UtilityMode._wpilib_opmode_metadata.name == "Arm Test" + assert UtilityMode._wpilib_opmode_metadata.text_color == Color.WHITE + assert UtilityMode._wpilib_opmode_metadata.background_color == Color.BLACK + assert wpilib.autonomous is autonomous + assert wpilib.teleop is teleop + assert wpilib.utility is utility + + +def test_opmode_decorators_expose_typed_overloads_and_docstrings(): + for decorator in (autonomous, teleop, utility): + overloads = get_overloads(decorator) + assert len(overloads) == 2 + for overload in overloads: + hints = get_type_hints(overload) + assert hints["name"] is str + assert hints["group"] is str + assert hints["description"] is str + assert hints["text_color"] == Color | None + assert hints["background_color"] == Color | None + + bare_return = get_type_hints(overloads[0])["return"] + configured_return = get_type_hints(overloads[1])["return"] + subtype = get_args(bare_return)[0] + configured_parameters, configured_result = get_args(configured_return) + assert get_origin(bare_return) is type + assert subtype.__bound__ is OpMode + assert get_args(configured_parameters[0])[0] is subtype + assert get_args(configured_result)[0] is subtype + assert decorator.__doc__ + assert "bare" in decorator.__doc__ + assert "configured" in decorator.__doc__ + + +def test_opmode_decorators_preserve_distinct_function_local_classes(): + def make_mode(name): + @utility(name=name) + class LocalMode(OpMode): + pass + + return LocalMode + + first_mode = make_mode("First Local") + second_mode = make_mode("Second Local") + + class MinimalRobot(OpModeRobot): + pass + + MinimalRobot() + + assert opmode_impl.decorated_opmodes() == (first_mode, second_mode) + assert sorted( + option.name for option in wsim.DriverStationSim.get_opmode_options() + ) == ["First Local", "Second Local"] + + +def test_opmode_decorator_rejects_invalid_class_and_duplicate_mode(): + with pytest.raises(TypeError, match="OpMode subclass"): + autonomous(type("NotAnOpMode", (), {})) + + @teleop + class DriveMode(OpMode): + pass + + with pytest.raises(ValueError, match="multiple opmode decorators"): + autonomous(DriveMode) + + +def test_opmode_detector_recognizes_imported_module_alias(): + source = """ + from wpilib import PeriodicOpMode + from wpilib import opmodes as modes + + @modes.utility + class UtilityMode(PeriodicOpMode): + pass + """ + + assert opmode_impl._has_opmode_decorator(textwrap.dedent(source), "utility.py") + + +@pytest.mark.parametrize("decorator", ["autonomous", "teleop", "utility"]) +@pytest.mark.parametrize("configured", [False, True]) +@pytest.mark.parametrize( + "import_statement, decorator_prefix", + [ + ("import wpilib.opmodes", "wpilib"), + ("import wpilib.opmodes", "wpilib.opmodes"), + ("import wpilib.opmodes as modes", "modes"), + ], +) +def test_opmode_robot_discovers_decorators_after_dotted_import( + monkeypatch, tmp_path, decorator, configured, import_statement, decorator_prefix +): + decorator_args = '(group="Drive")' if configured else "" + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/dotted_import.py": f""" + from wpilib import OpMode + {import_statement} + + @{decorator_prefix}.{decorator}{decorator_args} + class DiscoveredMode(OpMode): + pass + """, + }, + ) + assert "samplebot.opmodes.dotted_import" not in sys.modules + + robot_module.Robot() + + options = wsim.DriverStationSim.get_opmode_options() + assert [option.name for option in options] == ["DiscoveredMode"] + assert options[0].group == ("Drive" if configured else "") + assert "samplebot.opmodes.dotted_import" in sys.modules + + +@pytest.mark.parametrize("constructor_args", ["", "auto_discover=True"]) +def test_opmode_robot_auto_discovers_bounded_opmodes( + monkeypatch, tmp_path, constructor_args +): + pkg, robot_module = import_robot_package( + monkeypatch, + tmp_path, + f""" + import wpilib as wpi + + class Robot(wpi.OpModeRobot): + def __init__(self): + super().__init__({constructor_args}) + """, + { + "opmodes/auto_mode.py": """ + import wpilib as wpi + + @wpi.autonomous(name="Auto") + class AutoMode(wpi.PeriodicOpMode): + pass + """, + "opmodes/nested/__init__.py": "", + "opmodes/nested/tele_mode.py": """ + from wpilib import PeriodicOpMode + from wpilib.opmodes import teleop as tele + + @tele(group="Drive") + class TeleMode(PeriodicOpMode): + pass + """, + "opmodes/ignored.py": """ + from pathlib import Path + Path(__file__).with_name("ignored-imported").touch() + """, + }, + ) + + robot_module.Robot() + + options = { + option.name: option for option in wsim.DriverStationSim.get_opmode_options() + } + assert set(options) == {"Auto", "TeleMode"} + assert options["TeleMode"].group == "Drive" + assert not (pkg / "opmodes" / "ignored-imported").exists() + + +def test_opmode_robot_can_disable_auto_discovery(monkeypatch, tmp_path): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + @wpilib.autonomous + class ImportedMode(wpilib.OpMode): + pass + + class Robot(wpilib.OpModeRobot): + def __init__(self): + super().__init__(auto_discover=False) + """, + { + "opmodes/discovered_mode.py": """ + from pathlib import Path + import wpilib + + Path(__file__).with_name("mode-imported").touch() + + @wpilib.teleop + class DiscoveredMode(wpilib.OpMode): + pass + """, + }, + ) + + robot = robot_module.Robot() + + assert "samplebot.opmodes.discovered_mode" not in sys.modules + assert not (package_dir / "opmodes" / "mode-imported").exists() + assert not wsim.DriverStationSim.get_opmode_options() + + # Explicit publication must not reveal any automatically registered modes. + robot.publish_opmodes() + assert not wsim.DriverStationSim.get_opmode_options() + + robot.add_opmode(robot_module.ImportedMode, RobotMode.AUTONOMOUS, "Manual") + assert not wsim.DriverStationSim.get_opmode_options() + robot.publish_opmodes() + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "Manual" + ] + + +def test_disabled_auto_discovery_does_not_publish_pending_opmodes(): + class ManualMode(OpMode): + pass + + publisher = OpModeRobot() + publisher.add_opmode(ManualMode, RobotMode.AUTONOMOUS, "Manual") + assert not wsim.DriverStationSim.get_opmode_options() + + robot = OpModeRobot(auto_discover=False) + + assert not wsim.DriverStationSim.get_opmode_options() + publisher.publish_opmodes() + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "Manual" + ] + + +@pytest.mark.parametrize("has_init", [True, False]) +@pytest.mark.parametrize("preload", [True, False]) +def test_opmode_robot_discovers_multiple_implicit_namespace_modules( + monkeypatch, tmp_path, has_init, preload +): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/namespace/first.py": """ + import wpilib + + @wpilib.autonomous + class FirstMode(wpilib.OpMode): + pass + """, + "opmodes/namespace/second.py": """ + import wpilib + + @wpilib.teleop + class SecondMode(wpilib.OpMode): + pass + """, + }, + ) + + if not has_init: + (package_dir / "opmodes" / "__init__.py").unlink() + if preload: + importlib.import_module("samplebot.opmodes") + + robot_module.Robot() + + assert {option.name for option in wsim.DriverStationSim.get_opmode_options()} == { + "FirstMode", + "SecondMode", + } + + +@pytest.mark.parametrize("has_init", [True, False]) +@pytest.mark.parametrize( + ("robot_path", "robot_module_name", "opmodes_path", "expected_name"), + [ + ("samplebot/__init__.py", "samplebot", "samplebot/opmodes", "PackageMode"), + ( + "samplebot/robot/__init__.py", + "samplebot.robot", + "samplebot/robot/opmodes", + "NestedPackageMode", + ), + ("robot.py", "robot", "opmodes", "TopLevelMode"), + ], +) +def test_opmode_robot_derives_opmodes_package_from_robot_source( + monkeypatch, + tmp_path, + robot_path, + robot_module_name, + opmodes_path, + expected_name, + has_init, +): + robot_file = tmp_path / robot_path + robot_file.parent.mkdir(parents=True, exist_ok=True) + if robot_path == "samplebot/robot/__init__.py": + (tmp_path / "samplebot" / "__init__.py").write_text("") + robot_file.write_text( + "import wpilib\n\nclass Robot(wpilib.OpModeRobot):\n pass\n" + ) + opmodes_dir = tmp_path / opmodes_path + opmodes_dir.mkdir(parents=True) + if has_init: + (opmodes_dir / "__init__.py").write_text("") + (opmodes_dir / "mode.py").write_text( + "import wpilib\n\n" + f"@wpilib.autonomous(name={expected_name!r})\n" + "class Mode(wpilib.OpMode):\n pass\n" + ) + + monkeypatch.syspath_prepend(str(tmp_path)) + importlib.invalidate_caches() + robot_module = importlib.import_module(robot_module_name) + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + expected_name + ] + + +@pytest.mark.parametrize("has_init", [True, False]) +def test_opmode_robot_rejects_preloaded_package_origin_collision( + monkeypatch, tmp_path, caplog, has_init +): + conflict_dir = tmp_path / "conflict" + conflict_opmodes = conflict_dir / "opmodes" + conflict_opmodes.mkdir(parents=True) + (conflict_opmodes / "__init__.py").write_text("") + (conflict_opmodes / "mode.py").write_text( + "from pathlib import Path\n" + 'Path(__file__).with_name("conflicting-mode-imported").touch()\n' + ) + monkeypatch.syspath_prepend(str(conflict_dir)) + importlib.import_module("opmodes") + + robot_dir = tmp_path / "robot_project" + robot_dir.mkdir() + (robot_dir / "robot.py").write_text( + "import wpilib\n\nclass Robot(wpilib.OpModeRobot):\n pass\n" + ) + expected_opmodes = robot_dir / "opmodes" + expected_opmodes.mkdir() + if has_init: + (expected_opmodes / "__init__.py").write_text("") + (expected_opmodes / "mode.py").write_text( + "import wpilib\n\n" + "@wpilib.autonomous\n" + "class ExpectedMode(wpilib.OpMode):\n pass\n" + ) + monkeypatch.syspath_prepend(str(robot_dir)) + importlib.invalidate_caches() + + importlib.import_module("robot").Robot() + + assert not wsim.DriverStationSim.get_opmode_options() + assert not (conflict_opmodes / "conflicting-mode-imported").exists() + assert "opmodes.mode" in caplog.text + assert str(expected_opmodes / "mode.py") in caplog.text + + +def test_post_resolution_origin_mismatch_does_not_register_foreign_class( + monkeypatch, tmp_path, caplog +): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/raced_mode.py": """ + import wpilib + + @wpilib.autonomous(name="Expected") + class ExpectedMode(wpilib.OpMode): + pass + """, + }, + ) + foreign_path = tmp_path / "foreign_raced_mode.py" + foreign_path.write_text( + "import wpilib\n\n" + "@wpilib.teleop(name='Foreign')\n" + "class ForeignMode(wpilib.OpMode):\n pass\n" + ) + real_import_module = opmode_impl.importlib.import_module + + def raced_import(module_name): + if module_name != "samplebot.opmodes.raced_mode": + return real_import_module(module_name) + spec = importlib.util.spec_from_file_location(module_name, foreign_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + monkeypatch.setattr(opmode_impl, "_has_origin_collision", lambda *args: False) + monkeypatch.setattr(opmode_impl.importlib, "import_module", raced_import) + + robot_module.Robot() + + assert not wsim.DriverStationSim.get_opmode_options() + assert "samplebot.opmodes.raced_mode" in caplog.text + assert str(package_dir / "opmodes" / "raced_mode.py") in caplog.text + + +def test_rejected_scanned_origin_does_not_invalidate_explicit_class( + monkeypatch, tmp_path +): + conflict_dir = tmp_path / "conflict" + conflict_opmodes = conflict_dir / "opmodes" + conflict_opmodes.mkdir(parents=True) + (conflict_opmodes / "__init__.py").write_text("") + (conflict_opmodes / "base.py").write_text( + "import wpilib\n\n" + "@wpilib.autonomous(name='External Base')\n" + "class BaseMode(wpilib.OpMode):\n pass\n" + ) + monkeypatch.syspath_prepend(str(conflict_dir)) + importlib.import_module("opmodes.base") + + robot_dir = tmp_path / "robot_project" + robot_dir.mkdir() + (robot_dir / "robot.py").write_text( + "import wpilib\n\nclass Robot(wpilib.OpModeRobot):\n pass\n" + ) + expected_opmodes = robot_dir / "opmodes" + expected_opmodes.mkdir() + (expected_opmodes / "__init__.py").write_text("") + (expected_opmodes / "base.py").write_text( + "import wpilib\n\n" + "@wpilib.autonomous(name='Scanned Base')\n" + "class BaseMode(wpilib.OpMode):\n pass\n" + ) + (expected_opmodes / "child.py").write_text( + "from opmodes.base import BaseMode\n\n" "class ChildMode(BaseMode):\n pass\n" + ) + monkeypatch.syspath_prepend(str(robot_dir)) + importlib.invalidate_caches() + + importlib.import_module("robot").Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "External Base" + ] + + +def test_rejected_scanned_child_origin_does_not_invalidate_explicit_class( + monkeypatch, tmp_path +): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + import samplebot.external_base + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "external_base.py": """ + import wpilib + + @wpilib.autonomous(name="External Base") + class ExternalBaseMode(wpilib.OpMode): + pass + """, + "opmodes/subclass_only.py": """ + from pathlib import Path + from samplebot.external_base import ExternalBaseMode + + Path(__file__).with_name("colliding-child-imported").touch() + + class ExternalChildMode(ExternalBaseMode): + pass + """, + }, + ) + foreign_path = tmp_path / "foreign_subclass_only.py" + foreign_path.write_text("") + module_name = "samplebot.opmodes.subclass_only" + spec = importlib.util.spec_from_file_location(module_name, foreign_path) + foreign_module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = foreign_module + spec.loader.exec_module(foreign_module) + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "External Base" + ] + assert not (package_dir / "opmodes" / "colliding-child-imported").exists() + + +def test_automatic_publication_bypasses_python_override(): + @autonomous + class AutomaticMode(OpMode): + pass + + class OverrideRobot(OpModeRobot): + def publish_opmodes(self): + raise AssertionError("automatic publication called Python override") + + OverrideRobot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "AutomaticMode" + ] + + +def test_opmode_robot_registers_explicitly_imported_decorated_opmode( + monkeypatch, tmp_path +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib as wpi + import samplebot.external_mode + + class Robot(wpi.OpModeRobot): + def __init__(self): + super().__init__() + """, + { + "external_mode.py": """ + from wpilib import PeriodicOpMode + from wpilib.opmodes import teleop + + @teleop + class ImportedMode(PeriodicOpMode): + pass + """, + }, + ) + + robot_module.Robot() + + options = wsim.DriverStationSim.get_opmode_options() + assert [option.name for option in options] == ["ImportedMode"] + + +def test_opmode_robot_deduplicates_explicit_and_discovered_opmode( + monkeypatch, tmp_path +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + import samplebot.opmodes.overlap_mode + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/overlap_mode.py": """ + import wpilib + + @wpilib.teleop + class OverlapMode(wpilib.OpMode): + pass + """, + }, + ) + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "OverlapMode" + ] + + +def test_opmode_robot_prunes_failed_explicit_import(monkeypatch, tmp_path): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + try: + import samplebot.failed_mode + except RuntimeError: + pass + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "failed_mode.py": """ + import wpilib + + @wpilib.autonomous + class FailedMode(wpilib.OpMode): + pass + + raise RuntimeError("expected explicit import failure") + """, + }, + ) + + assert "samplebot.failed_mode" not in sys.modules + robot_module.Robot() + + assert not wsim.DriverStationSim.get_opmode_options() + + +def test_opmode_robot_replaces_reloaded_class_generation(monkeypatch, tmp_path): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + import samplebot.external_mode + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "external_mode.py": """ + import wpilib + + @wpilib.utility + class ReloadedMode(wpilib.OpMode): + pass + """, + }, + ) + mode_module = sys.modules["samplebot.external_mode"] + old_class = mode_module.ReloadedMode + importlib.reload(mode_module) + + robot_module.Robot() + + options = wsim.DriverStationSim.get_opmode_options() + assert [option.name for option in options] == ["ReloadedMode"] + assert opmode_impl.decorated_opmodes() == (mode_module.ReloadedMode,) + assert mode_module.ReloadedMode is not old_class + + +def test_opmode_robot_prunes_failed_reload_generation(monkeypatch, tmp_path): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + import samplebot.external_mode + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "external_mode.py": """ + import wpilib + + @wpilib.utility + class ReloadedMode(wpilib.OpMode): + pass + """, + }, + ) + mode_path = package_dir / "external_mode.py" + mode_path.write_text( + "import wpilib\n\n" + "class ReloadedMode(wpilib.OpMode):\n pass\n\n" + 'raise RuntimeError("expected reload failure")\n' + ) + for bytecode_path in (package_dir / "__pycache__").glob("external_mode.*.pyc"): + bytecode_path.unlink() + importlib.invalidate_caches() + mode_module = sys.modules["samplebot.external_mode"] + with pytest.raises(RuntimeError, match="expected reload failure"): + importlib.reload(mode_module) + + robot_module.Robot() + + assert not wsim.DriverStationSim.get_opmode_options() + assert opmode_impl.decorated_opmodes() == () + assert "_wpilib_opmode_metadata" not in mode_module.ReloadedMode.__dict__ + + +def test_opmode_robot_rejects_decorated_non_leaf(caplog): + @autonomous + class BaseMode(OpMode): + pass + + class MiddleMode(BaseMode): + pass + + class LeafMode(MiddleMode): + pass + + class MinimalRobot(OpModeRobot): + pass + + MinimalRobot() + + assert not wsim.DriverStationSim.get_opmode_options() + assert "BaseMode" in caplog.text + assert "MiddleMode" in caplog.text + + class CommonMode(OpMode): + pass + + @utility + class LeafUtilityMode(CommonMode): + pass + + MinimalRobot() + + options = wsim.DriverStationSim.get_opmode_options() + assert [option.name for option in options] == ["LeafUtilityMode"] + + +def test_opmode_robot_reports_invalid_metadata_per_class(monkeypatch, tmp_path, caplog): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/metadata_modes.py": """ + import wpilib + + @wpilib.autonomous(name=1) + class BadName(wpilib.OpMode): + pass + + @wpilib.teleop(group=None) + class BadGroup(wpilib.OpMode): + pass + + @wpilib.utility(description=[]) + class BadDescription(wpilib.OpMode): + pass + + @wpilib.autonomous(text_color="white") + class BadTextColor(wpilib.OpMode): + pass + + @wpilib.teleop(background_color=object()) + class BadBackgroundColor(wpilib.OpMode): + pass + + @wpilib.utility(name="Good") + class GoodMode(wpilib.OpMode): + pass + """, + }, + ) + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "Good" + ] + for class_name, field_name in ( + ("BadName", "name"), + ("BadGroup", "group"), + ("BadDescription", "description"), + ("BadTextColor", "text_color"), + ("BadBackgroundColor", "background_color"), + ): + assert f"samplebot.opmodes.metadata_modes.{class_name}" in caplog.text + assert field_name in caplog.text + + +@pytest.mark.parametrize("overridden_method", ["add_opmode", "add_opmode_factory"]) +@pytest.mark.parametrize("colored", [False, True]) +def test_opmode_robot_automatic_registration_bypasses_overrides( + monkeypatch, caplog, overridden_method, colored +): + @utility( + text_color=Color.WHITE if colored else None, + background_color=Color.BLACK if colored else None, + ) + class AutomaticMode(OpMode): + pass + + class ManualMode(OpMode): + pass + + class Robot(OpModeRobot): + def __init__(self): + super().__init__() + self.registered = [] + + original = getattr(OpModeRobot, overridden_method) + + def register(self, *args, **kwargs): + self.registered.append(args) + return original(self, *args, **kwargs) + + monkeypatch.setattr(Robot, overridden_method, register) + + robot = Robot() + + options = wsim.DriverStationSim.get_opmode_options() + assert [option.name for option in options] == ["AutomaticMode"] + assert options[0].text_color == (0xFFFFFF if colored else -1) + assert options[0].background_color == (0x000000 if colored else -1) + assert robot.registered == [] + assert "Could not register" not in caplog.text + + # Explicit registration still dispatches through the user's overrides. + robot.add_opmode(ManualMode, RobotMode.UTILITY, "ManualMode") + robot.publish_opmodes() + assert len(robot.registered) == 1 + assert sorted( + option.name for option in wsim.DriverStationSim.get_opmode_options() + ) == ["AutomaticMode", "ManualMode"] + + +def test_opmode_robot_continues_after_per_class_registration_failure( + monkeypatch, tmp_path, caplog +): + original = OpModeRobotBase.add_opmode_factory + + def register(robot, mode, name, *args): + if name == "BadMode": + raise RuntimeError("expected registration failure") + return original(robot, mode, name, *args) + + monkeypatch.setattr(OpModeRobotBase, "add_opmode_factory", register) + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/modes.py": """ + import wpilib + + @wpilib.autonomous + class BadMode(wpilib.OpMode): + pass + + @wpilib.teleop + class GoodMode(wpilib.OpMode): + pass + """, + }, + ) + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "GoodMode" + ] + assert "samplebot.opmodes.modes.BadMode" in caplog.text + assert "expected registration failure" in caplog.text + + +def test_opmode_robot_registers_explicit_base_with_unimported_scanned_subclass( + monkeypatch, tmp_path, caplog +): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + import samplebot.external_base + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "external_base.py": """ + import wpilib + + @wpilib.autonomous + class ExternalBaseMode(wpilib.OpMode): + pass + """, + "opmodes/subclass_only.py": """ + from pathlib import Path + from samplebot.external_base import ExternalBaseMode + + Path(__file__).with_name("external-subclass-imported").touch() + + class ExternalChildMode(ExternalBaseMode): + pass + """, + }, + ) + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "ExternalBaseMode" + ] + assert not (package_dir / "opmodes" / "external-subclass-imported").exists() + assert "must not be subclassed" not in caplog.text + + +@pytest.mark.parametrize("import_child", [False, True]) +@pytest.mark.parametrize( + "child_source, registers", + [ + pytest.param("Base = object\nclass Child(Base): pass", True, id="assignment"), + pytest.param( + "Base = Other = object\nclass Child(Base): pass", True, id="chained" + ), + pytest.param( + "Base, Other = object, object\nclass Child(Base): pass", + True, + id="unpacking", + ), + pytest.param( + "[Base, *Other] = [object, object]\nclass Child(Base): pass", + True, + id="starred-unpacking", + ), + pytest.param( + "Base: type = object\nclass Child(Base): pass", True, id="annotated" + ), + pytest.param( + "class Replacement:\n def __ror__(self, other): return object\n" + "Base |= Replacement()\nclass Child(Base): pass", + True, + id="augmented", + ), + pytest.param( + "(Base := object)\nclass Child(Base): pass", True, id="named-expression" + ), + pytest.param( + "class Container:\n Base = object\n class Child(Base): pass", + True, + id="class-scope-shadow", + ), + pytest.param( + "def make_child():\n Base = object\n class Child(Base): pass", + True, + id="function-scope-shadow", + ), + pytest.param( + "from types import SimpleNamespace\n" + "import samplebot.external_base as modes\n" + "modes = SimpleNamespace(Base=object)\nclass Child(modes.Base): pass", + True, + id="module-alias", + ), + pytest.param( + "for Base in [object]: pass\nclass Child(Base): pass", + True, + id="for-target", + ), + pytest.param( + "from contextlib import nullcontext\n" + "with nullcontext(object) as Base:\n class Child(Base): pass", + True, + id="with-target", + ), + pytest.param( + "def make_child(Base):\n class Child(Base): pass\n return Child\n" + "Child = make_child(object)", + True, + id="parameter", + ), + pytest.param( + "if False: Base = object\nclass Child(Base): pass", + False, + id="false-branch", + ), + pytest.param( + "try: pass\nexcept Exception: Base = object\nclass Child(Base): pass", + False, + id="exception-path", + ), + pytest.param("class Child(Base): pass", False, id="unchanged-import"), + pytest.param( + "class Child(Base): pass\nBase = object", False, id="later-assignment" + ), + pytest.param( + "Base: type\nclass Child(Base): pass", False, id="annotation-only" + ), + pytest.param( + "def unrelated():\n Base = object\nclass Child(Base): pass", + False, + id="unrelated-scope", + ), + pytest.param( + "unused = lambda: (Base := object)\nclass Child(Base): pass", + False, + id="unrelated-lambda-scope", + ), + pytest.param( + "unused = lambda arg=(Base := object): arg\nclass Child(Base): pass", + True, + id="lambda-default-in-outer-scope", + ), + pytest.param( + "Base = object\nfrom samplebot.external_base import Base\n" + "class Child(Base): pass", + False, + id="reimport", + ), + ], +) +def test_opmode_robot_checks_runtime_subclasses( + monkeypatch, tmp_path, caplog, child_source, registers, import_child +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + import samplebot.external_base + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "external_base.py": """ + import wpilib + + @wpilib.autonomous + class Base(wpilib.OpMode): + pass + """, + "opmodes/subclass_only.py": ( + "from samplebot.external_base import Base\n" + child_source + "\n" + ), + }, + ) + + if import_child: + importlib.import_module("samplebot.opmodes.subclass_only") + + robot_module.Robot() + + should_register = not import_child or registers + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == ( + ["Base"] if should_register else [] + ) + assert ("must not be subclassed" in caplog.text) is not should_register + assert ("samplebot.opmodes.subclass_only" in sys.modules) is import_child + + +@pytest.mark.parametrize("import_child", [False, True]) +@pytest.mark.parametrize("import_name", ["Base", "*"]) +def test_opmode_robot_checks_subclasses_only_in_imported_modules( + monkeypatch, tmp_path, caplog, import_child, import_name +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/base_mode.py": """ + import wpilib + + @wpilib.autonomous + class Base(wpilib.OpMode): + pass + """, + "opmodes/subclass_only.py": ( + f"from .base_mode import {import_name}\nclass Child(Base): pass\n" + ), + }, + ) + if import_child: + importlib.import_module("samplebot.opmodes.subclass_only") + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == ( + [] if import_child else ["Base"] + ) + assert ("must not be subclassed" in caplog.text) is import_child + assert ("samplebot.opmodes.subclass_only" in sys.modules) is import_child + + +def test_opmode_robot_continues_after_candidate_import_failure( + monkeypatch, tmp_path, caplog +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib as wpi + + class Robot(wpi.OpModeRobot): + pass + """, + { + "opmodes/bad_mode.py": """ + from wpilib import PeriodicOpMode, teleop + + @teleop + class BadMode(PeriodicOpMode): + pass + + raise RuntimeError("expected candidate import failure") + """, + "opmodes/good_mode.py": """ + from wpilib import PeriodicOpMode, utility + + @utility + class GoodMode(PeriodicOpMode): + pass + """, + }, + ) + + robot_module.Robot() + + options = wsim.DriverStationSim.get_opmode_options() + assert {option.name for option in options} == {"GoodMode"} + assert "bad_mode" in caplog.text + assert "expected candidate import failure" in caplog.text + + +def test_opmode_robot_retains_transitively_imported_candidate_after_failure( + monkeypatch, tmp_path, caplog +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib as wpi + + class Robot(wpi.OpModeRobot): + pass + """, + { + "opmodes/bad_mode.py": """ + from samplebot.opmodes import good_mode + from wpilib import PeriodicOpMode, teleop + + @teleop + class BadMode(PeriodicOpMode): + pass + + raise RuntimeError("expected candidate import failure") + """, + "opmodes/good_mode.py": """ + from wpilib import PeriodicOpMode, utility + + @utility + class GoodMode(PeriodicOpMode): + pass + """, + }, + ) + + robot_module.Robot() + + options = wsim.DriverStationSim.get_opmode_options() + assert {option.name for option in options} == {"GoodMode"} + assert "bad_mode" in caplog.text + assert "expected candidate import failure" in caplog.text + + +def test_opmode_robot_rolls_back_nested_failed_candidates( + monkeypatch, tmp_path, caplog +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib as wpi + + class Robot(wpi.OpModeRobot): + pass + """, + { + "opmodes/bad_a.py": """ + from samplebot.opmodes import bad_b + from wpilib import PeriodicOpMode, teleop + + @teleop + class BadAMode(PeriodicOpMode): + pass + """, + "opmodes/bad_b.py": """ + from wpilib import PeriodicOpMode, utility + + @utility + class BadBMode(PeriodicOpMode): + pass + + raise RuntimeError("expected nested candidate import failure") + """, + }, + ) + + robot_module.Robot() + + assert not wsim.DriverStationSim.get_opmode_options() + assert "bad_a" in caplog.text + assert "Could not import OpMode module samplebot.opmodes.bad_b" in caplog.text + assert "expected nested candidate import failure" in caplog.text + + +def test_opmode_robot_discovers_encoded_python_source(monkeypatch, tmp_path): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + {"opmodes/encoded_mode.py": ""}, + ) + (package_dir / "opmodes" / "encoded_mode.py").write_bytes( + b"# -*- coding: latin-1 -*-\n" + b"import wpilib\n\n" + b"@wpilib.autonomous(name='Caf\xe9')\n" + b"class EncodedMode(wpilib.OpMode):\n pass\n" + ) + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "Caf\N{LATIN SMALL LETTER E WITH ACUTE}" + ] + + +def test_opmode_robot_continues_after_source_decode_failure( + monkeypatch, tmp_path, caplog +): + package_dir, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/bad_encoding.py": "", + "opmodes/good_mode.py": """ + import wpilib + + @wpilib.utility + class GoodMode(wpilib.OpMode): + pass + """, + }, + ) + (package_dir / "opmodes" / "bad_encoding.py").write_bytes( + b"# coding: utf-8\n\xff\n" + ) + + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "GoodMode" + ] + assert "bad_encoding.py" in caplog.text + + +def test_opmode_robot_continues_after_source_read_failure( + monkeypatch, tmp_path, caplog +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib + + class Robot(wpilib.OpModeRobot): + pass + """, + { + "opmodes/unreadable.py": "", + "opmodes/good_mode.py": """ + import wpilib + + @wpilib.utility + class GoodMode(wpilib.OpMode): + pass + """, + }, + ) + real_open = opmode_impl.tokenize.open + + def open_source(filename): + if Path(filename).name == "unreadable.py": + raise OSError("expected source read failure") + return real_open(filename) + + monkeypatch.setattr(opmode_impl.tokenize, "open", open_source) + robot_module.Robot() + + assert [option.name for option in wsim.DriverStationSim.get_opmode_options()] == [ + "GoodMode" + ] + assert "unreadable.py" in caplog.text + assert "expected source read failure" in caplog.text + + +def test_opmode_robot_continues_after_candidate_parse_failure( + monkeypatch, tmp_path, caplog +): + _, robot_module = import_robot_package( + monkeypatch, + tmp_path, + """ + import wpilib as wpi + + class Robot(wpi.OpModeRobot): + pass + """, + { + "opmodes/bad_syntax.py": """ + from wpilib import PeriodicOpMode, teleop + + @teleop + class BrokenMode(PeriodicOpMode) + pass + """, + "opmodes/good_mode.py": """ + from wpilib import PeriodicOpMode, utility + + @utility + class GoodMode(PeriodicOpMode): + pass + """, + }, + ) + + robot_module.Robot() + + options = wsim.DriverStationSim.get_opmode_options() + assert {option.name for option in options} == {"GoodMode"} + assert "bad_syntax.py" in caplog.text + assert "expected ':'" in caplog.text + + class MockOpMode(OpMode): def __init__(self): super().__init__() diff --git a/subprojects/robotpy-wpilib/wpilib/__init__.py b/subprojects/robotpy-wpilib/wpilib/__init__.py index bec259c30..92e9af9b6 100644 --- a/subprojects/robotpy-wpilib/wpilib/__init__.py +++ b/subprojects/robotpy-wpilib/wpilib/__init__.py @@ -242,9 +242,9 @@ del _init__wpilib -from .opmoderobot import OpModeRobot +from .opmodes import OpModeRobot, autonomous, teleop, utility -__all__ += ["OpModeRobot"] +__all__ += ["OpModeRobot", "autonomous", "teleop", "utility"] from .cameraserver import CameraServer from .deployinfo import get_deploy_data diff --git a/subprojects/robotpy-wpilib/wpilib/_impl/opmode.py b/subprojects/robotpy-wpilib/wpilib/_impl/opmode.py new file mode 100644 index 000000000..cf05c1923 --- /dev/null +++ b/subprojects/robotpy-wpilib/wpilib/_impl/opmode.py @@ -0,0 +1,427 @@ +from __future__ import annotations + +import ast +from collections.abc import Callable +from dataclasses import dataclass +from functools import partial +import importlib +from importlib.machinery import PathFinder +import inspect +from pathlib import Path +import sys +import tokenize +from typing import Any, TypeVar + +from hal import RobotMode +from wpiutil import Color +from wpilib._wpilib import OpMode, OpModeRobotBase, RobotState + +from .report_error import report_error + + +@dataclass(frozen=True) +class OpModeMetadata: + mode: RobotMode + name: str + group: str + description: str + text_color: Color | None + background_color: Color | None + + +_decorated_opmodes: list[type[OpMode]] = [] +_OpModeT = TypeVar("_OpModeT", bound=OpMode) + + +def attach_metadata( + cls: type[_OpModeT], + *, + mode: RobotMode, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> type[_OpModeT]: + if not inspect.isclass(cls) or not issubclass(cls, OpMode): + raise TypeError("opmode decorator must be applied to an OpMode subclass") + if "_wpilib_opmode_metadata" in cls.__dict__: + raise ValueError("multiple opmode decorators are not allowed") + + cls._wpilib_opmode_metadata = OpModeMetadata( + mode, + cls.__name__ if name == "" else name, + group, + description, + text_color, + background_color, + ) + identity = (cls.__module__, cls.__qualname__) + if "" not in cls.__qualname__: + _decorated_opmodes[:] = [ + decorated_cls + for decorated_cls in _decorated_opmodes + if (decorated_cls.__module__, decorated_cls.__qualname__) != identity + ] + _decorated_opmodes.append(cls) + return cls + + +def _is_current_class_generation(cls: type[OpMode]) -> bool: + module = sys.modules.get(cls.__module__) + if module is None: + return False + parts = cls.__qualname__.split(".") + if "" in parts: + return True + + current = module.__dict__.get(parts[0]) + for part in parts[1:]: + namespace = getattr(current, "__dict__", None) + if namespace is None: + return False + current = namespace.get(part) + return current is cls + + +def decorated_opmodes() -> tuple[type[OpMode], ...]: + _decorated_opmodes[:] = [ + cls for cls in _decorated_opmodes if _is_current_class_generation(cls) + ] + return tuple(_decorated_opmodes) + + +def _tree_has_opmode_decorator(tree: ast.Module) -> bool: + decorator_names = {"autonomous", "teleop", "utility"} + imported_decorators: set[str] = set() + module_aliases: set[tuple[str, ...]] = set() + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for imported in node.names: + if imported.name not in {"wpilib", "wpilib.opmodes"}: + continue + if imported.asname is not None: + module_aliases.add((imported.asname,)) + else: + module_aliases.add(tuple(imported.name.split("."))) + # An unaliased dotted import also binds the root package. + module_aliases.add((imported.name.split(".")[0],)) + elif isinstance(node, ast.ImportFrom) and node.module in { + "wpilib", + "wpilib.opmodes", + }: + for imported in node.names: + if node.module == "wpilib" and imported.name == "opmodes": + module_aliases.add((imported.asname or imported.name,)) + elif imported.name == "*": + imported_decorators.update(decorator_names) + elif imported.name in decorator_names: + imported_decorators.add(imported.asname or imported.name) + + def attribute_path(node: ast.expr) -> tuple[str, ...] | None: + parts: list[str] = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + if not isinstance(node, ast.Name): + return None + parts.append(node.id) + return tuple(reversed(parts)) + + for node in ast.walk(tree): + if not isinstance(node, ast.ClassDef): + continue + for decorator in node.decorator_list: + if isinstance(decorator, ast.Call): + decorator = decorator.func + if isinstance(decorator, ast.Name): + if decorator.id in imported_decorators: + return True + elif isinstance(decorator, ast.Attribute): + path = attribute_path(decorator) + if ( + path is not None + and path[-1] in decorator_names + and path[:-1] in module_aliases + ): + return True + return False + + +def _has_opmode_decorator(source: str, filename: str) -> bool: + return _tree_has_opmode_decorator(ast.parse(source, filename=filename)) + + +@dataclass(frozen=True) +class _ScannedModule: + name: str + path: Path + tree: ast.Module + + +def _scan_modules(package_dir: Path, package_name: str) -> list[_ScannedModule]: + modules: list[_ScannedModule] = [] + for source_path in sorted(package_dir.rglob("*.py")): + try: + with tokenize.open(source_path) as source_file: + source = source_file.read() + tree = ast.parse(source, filename=str(source_path)) + except (OSError, UnicodeError, SyntaxError) as exc: + report_error(f"Could not read or parse OpMode module {source_path}: {exc}") + continue + + relative_path = source_path.relative_to(package_dir) + is_package = relative_path.name == "__init__.py" + module_parts = ( + relative_path.parent.parts + if is_package + else relative_path.with_suffix("").parts + ) + modules.append( + _ScannedModule( + ".".join((package_name, *module_parts)), + source_path, + tree, + ) + ) + return modules + + +def _path_from_origin(origin: str | None) -> Path | None: + if origin is None or origin in {"built-in", "frozen"}: + return None + return Path(origin).resolve() + + +def _loaded_origin(module_name: str) -> Path | None: + module = sys.modules.get(module_name) + if module is None: + return None + origin = getattr(getattr(module, "__spec__", None), "origin", None) + return _path_from_origin(origin or getattr(module, "__file__", None)) + + +def _resolved_origin(module_name: str) -> Path | None: + search_path = None + spec = None + parts = module_name.split(".") + for index in range(len(parts)): + current_name = ".".join(parts[: index + 1]) + loaded = sys.modules.get(current_name) + if loaded is not None: + spec = getattr(loaded, "__spec__", None) + else: + try: + spec = PathFinder.find_spec(current_name, search_path) + except KeyError: + namespace_paths = [ + str(Path(entry) / parts[index]) + for entry in (search_path or sys.path) + if (Path(entry) / parts[index]).is_dir() + ] + if not namespace_paths or index == len(parts) - 1: + return None + search_path = namespace_paths + continue + if spec is None: + return None + search_locations = spec.submodule_search_locations + if search_locations is None: + search_path = None + else: + search_path = list(getattr(search_locations, "_path", search_locations)) + return _path_from_origin(spec.origin) + + +def _expected_origins( + scanned_module: _ScannedModule, package_dir: Path, package_name: str +) -> dict[str, Path]: + expected: dict[str, Path] = {} + relative = scanned_module.path.relative_to(package_dir) + directories = relative.parent.parts + for index in range(len(directories) + 1): + package_path = package_dir.joinpath(*directories[:index]) + package_init = package_path / "__init__.py" + if package_init.is_file(): + module_name = ".".join((package_name, *directories[:index])) + expected[module_name] = package_init.resolve() + expected[scanned_module.name] = scanned_module.path.resolve() + return expected + + +def _has_origin_collision( + scanned_module: _ScannedModule, package_dir: Path, package_name: str +) -> bool: + expected_origins = _expected_origins(scanned_module, package_dir, package_name) + for module_name, expected_path in expected_origins.items(): + if module_name not in sys.modules: + continue + actual_path = _loaded_origin(module_name) + if actual_path != expected_path: + report_error( + f"Refusing to import OpMode module {scanned_module.name} from " + f"scanned path {scanned_module.path}: preloaded module " + f"{module_name} has origin {actual_path}" + ) + return True + + resolved_path = _resolved_origin(scanned_module.name) + if resolved_path is not None and resolved_path != scanned_module.path.resolve(): + report_error( + f"Refusing to import OpMode module {scanned_module.name} from scanned " + f"path {scanned_module.path}: module resolves to {resolved_path}" + ) + return True + return False + + +def _import_candidates( + modules: list[_ScannedModule], package_dir: Path, package_name: str +) -> None: + for scanned_module in modules: + if not _tree_has_opmode_decorator(scanned_module.tree): + continue + module_name = scanned_module.name + if _has_origin_collision(scanned_module, package_dir, package_name): + continue + registry_before = list(_decorated_opmodes) + origin_mismatch = False + try: + imported_module = importlib.import_module(module_name) + if _loaded_origin(module_name) != scanned_module.path.resolve(): + origin_mismatch = True + raise ImportError( + f"imported origin {getattr(imported_module, '__file__', None)} " + f"does not match scanned path {scanned_module.path}" + ) + except Exception as exc: + new_loaded_classes = [ + cls + for cls in _decorated_opmodes + if cls not in registry_before + and cls.__module__ in sys.modules + and not (origin_mismatch and cls.__module__ == module_name) + ] + _decorated_opmodes[:] = registry_before + new_loaded_classes + report_error( + f"Could not import OpMode module {module_name}: {exc}", + print_trace=True, + ) + + +def register_opmode( + robot: OpModeRobotBase, + opmode_cls: type, + mode: RobotMode, + name: str, + group: str | None = None, + description: str | None = None, + text_color: Color | None = None, + background_color: Color | None = None, + *, + add_factory: Callable[..., None] | None = None, +) -> None: + # Discovery runs before the subclass constructor has finished. Use the + # native registration method unless an explicit manual call supplies its + # own factory registration method. + if add_factory is None: + add_factory = partial(OpModeRobotBase.add_opmode_factory, robot) + + def make_opmode_instance() -> OpMode: + # Try to instantiate with robot argument first. + try: + return opmode_cls(robot) + except TypeError: + return opmode_cls() + + if text_color is None or background_color is None: + add_factory(mode, name, group or "", description or "", make_opmode_instance) + else: + add_factory( + mode, + name, + group or "", + description or "", + text_color, + background_color, + make_opmode_instance, + ) + + +def _find_subclass(cls: type[OpMode]) -> str | None: + subclasses = cls.__subclasses__() + if subclasses: + subclass = subclasses[0] + return f"{subclass.__module__}.{subclass.__qualname__}" + return None + + +def _invalid_metadata(metadata: OpModeMetadata) -> str | None: + for field_name in ("name", "group", "description"): + if not isinstance(getattr(metadata, field_name), str): + return f"{field_name} must be a string" + for field_name in ("text_color", "background_color"): + value = getattr(metadata, field_name) + if value is not None and not isinstance(value, Color): + return f"{field_name} must be a wpiutil.Color or None" + return None + + +def discover_and_register(robot: OpModeRobotBase) -> None: + robot_source = Path(inspect.getfile(type(robot))) + package_dir = robot_source.parent / "opmodes" + if package_dir.is_dir(): + robot_module = type(robot).__module__ + robot_package = ( + robot_module + if robot_source.name == "__init__.py" + else robot_module.rpartition(".")[0] + ) + package_name = f"{robot_package}.opmodes" if robot_package else "opmodes" + scanned_modules = _scan_modules(package_dir, package_name) + _import_candidates(scanned_modules, package_dir, package_name) + + registered: set[tuple[type[OpMode], RobotMode]] = set() + registrations: list[tuple[type[OpMode], OpModeMetadata]] = [] + for cls in decorated_opmodes(): + metadata = cls.__dict__.get("_wpilib_opmode_metadata") + if metadata is None or not issubclass(cls, OpMode): + continue + key = (cls, metadata.mode) + if key in registered: + continue + registered.add(key) + class_name = f"{cls.__module__}.{cls.__qualname__}" + metadata_error = _invalid_metadata(metadata) + if metadata_error is not None: + report_error(f"Invalid OpMode metadata for {class_name}: {metadata_error}") + continue + subclass = _find_subclass(cls) + if subclass is not None: + report_error( + f"Decorated OpMode {cls.__module__}.{cls.__qualname__} must not be " + f"subclassed; found subclass {subclass}" + ) + continue + registrations.append((cls, metadata)) + + for cls, metadata in registrations: + try: + register_opmode( + robot, + cls, + metadata.mode, + metadata.name, + metadata.group, + metadata.description, + metadata.text_color, + metadata.background_color, + ) + except Exception as exc: + report_error( + f"Could not register decorated OpMode " + f"{cls.__module__}.{cls.__qualname__}: {exc}", + print_trace=True, + ) + RobotState.publish_opmodes() diff --git a/subprojects/robotpy-wpilib/wpilib/opmoderobot.py b/subprojects/robotpy-wpilib/wpilib/opmoderobot.py deleted file mode 100644 index 746fde6ba..000000000 --- a/subprojects/robotpy-wpilib/wpilib/opmoderobot.py +++ /dev/null @@ -1,76 +0,0 @@ -from hal import RobotMode -from typing import Optional -from wpiutil import Color - -__all__ = ["OpModeRobot"] - -from ._wpilib import OpModeRobotBase, OpMode - - -class OpModeRobot(OpModeRobotBase): - """ - OpModeRobot implements the opmode-based robot program framework. - - The OpModeRobot class is intended to be subclassed by a user creating a robot - program. - - Opmodes are constructed when selected on the driver station, and destroyed - when the robot is disabled after being enabled or a different opmode is - selected. When no opmode is selected, none_periodic() is called. The - driver_station_connected() function is called the first time the driver station - connects to the robot. - """ - - def __init__(self): - super().__init__() - - def add_opmode( - self, - opmode_cls: type, - mode: RobotMode, - name: str, - group: Optional[str] = None, - description: Optional[str] = None, - text_color: Optional[Color] = None, - background_color: Optional[Color] = None, - ) -> None: - """ - Adds an operating mode option. It's necessary to call publish_opmodes() to - make the added modes visible to the driver station. - - The text_color and background_color parameters are optional, but setting - only one has no effect (if only one is provided, it will be ignored). - - :param opmode_cls: opmode class; must be a public, non-abstract subclass of OpMode - with a constructor that either takes no arguments or accepts a - single argument of this class's type (the latter is preferred). - :param mode: robot mode - :param name: name of the operating mode - :param group: group of the operating mode - :param description: description of the operating mode - :param text_color: text color - :param background_color: background color - """ - - def make_opmode_instance() -> OpMode: - # Try to instantiate with robot argument first - try: - return opmode_cls(self) # type: ignore - except TypeError: - # Fallback to no-argument constructor - return opmode_cls() # type: ignore - - if text_color is None or background_color is None: - self.add_opmode_factory( - mode, name, group or "", description or "", make_opmode_instance - ) - else: - self.add_opmode_factory( - mode, - name, - group or "", - description or "", - text_color, - background_color, - make_opmode_instance, - ) diff --git a/subprojects/robotpy-wpilib/wpilib/opmodes.py b/subprojects/robotpy-wpilib/wpilib/opmodes.py new file mode 100644 index 000000000..babe3a1e7 --- /dev/null +++ b/subprojects/robotpy-wpilib/wpilib/opmodes.py @@ -0,0 +1,262 @@ +from collections.abc import Callable +from typing import Optional, TypeVar, overload + +from hal import RobotMode +from wpiutil import Color + +__all__ = ["OpModeRobot", "autonomous", "teleop", "utility"] + +from ._impl import opmode as _opmode +from ._wpilib import OpModeRobotBase, OpMode + +_OpModeT = TypeVar("_OpModeT", bound=OpMode) + + +def _apply_opmode_decorator( + cls: type[_OpModeT] | None, + *, + mode: RobotMode, + name: str, + group: str, + description: str, + text_color: Color | None, + background_color: Color | None, +) -> type[_OpModeT] | Callable[[type[_OpModeT]], type[_OpModeT]]: + def apply(opmode_cls: type[_OpModeT]) -> type[_OpModeT]: + return _opmode.attach_metadata( + opmode_cls, + mode=mode, + name=name, + group=group, + description=description, + text_color=text_color, + background_color=background_color, + ) + + return apply if cls is None else apply(cls) + + +@overload +def autonomous( + cls: type[_OpModeT], + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> type[_OpModeT]: ... + + +@overload +def autonomous( + cls: None = None, + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> Callable[[type[_OpModeT]], type[_OpModeT]]: ... + + +def autonomous( + cls: type[_OpModeT] | None = None, + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> type[_OpModeT] | Callable[[type[_OpModeT]], type[_OpModeT]]: + """Mark an OpMode subclass for autonomous automatic registration. + + Use this decorator bare (``@autonomous``) or configured + (``@autonomous(name=..., group=...)``). The optional description and colors + are published with the Driver Station option. + """ + return _apply_opmode_decorator( + cls, + mode=RobotMode.AUTONOMOUS, + name=name, + group=group, + description=description, + text_color=text_color, + background_color=background_color, + ) + + +@overload +def teleop( + cls: type[_OpModeT], + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> type[_OpModeT]: ... + + +@overload +def teleop( + cls: None = None, + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> Callable[[type[_OpModeT]], type[_OpModeT]]: ... + + +def teleop( + cls: type[_OpModeT] | None = None, + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> type[_OpModeT] | Callable[[type[_OpModeT]], type[_OpModeT]]: + """Mark an OpMode subclass for teleoperated automatic registration. + + Use this decorator bare (``@teleop``) or configured + (``@teleop(name=..., group=...)``). The optional description and colors are + published with the Driver Station option. + """ + return _apply_opmode_decorator( + cls, + mode=RobotMode.TELEOPERATED, + name=name, + group=group, + description=description, + text_color=text_color, + background_color=background_color, + ) + + +@overload +def utility( + cls: type[_OpModeT], + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> type[_OpModeT]: ... + + +@overload +def utility( + cls: None = None, + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> Callable[[type[_OpModeT]], type[_OpModeT]]: ... + + +def utility( + cls: type[_OpModeT] | None = None, + *, + name: str = "", + group: str = "", + description: str = "", + text_color: Color | None = None, + background_color: Color | None = None, +) -> type[_OpModeT] | Callable[[type[_OpModeT]], type[_OpModeT]]: + """Mark an OpMode subclass for utility automatic registration. + + Use this decorator bare (``@utility``) or configured + (``@utility(name=..., group=...)``). The optional description and colors are + published with the Driver Station option. + """ + return _apply_opmode_decorator( + cls, + mode=RobotMode.UTILITY, + name=name, + group=group, + description=description, + text_color=text_color, + background_color=background_color, + ) + + +class OpModeRobot(OpModeRobotBase): + """ + OpModeRobot implements the opmode-based robot program framework. + + Base class for a robot program that uses selectable operating modes (OpModes). + + Create your robot class by inheriting from OpModeRobot. Mark your OpMode + classes with ``@autonomous``, ``@teleop``, or ``@utility``. OpModeRobot + automatically registers these classes when you import them or place them + in an ``opmodes`` package next to your robot module. + + To disable automatic discovery and registration, call + ``super().__init__(auto_discover=False)`` in your robot constructor. Use + ``add_opmode()`` and ``publish_opmodes()`` to register and publish modes manually. + + Selecting an OpMode on the Driver Station creates a new instance of it. + That instance is discarded when you select a different OpMode or disable + the robot after enabling it. + + Override ``none_periodic()`` to run code repeatedly while no OpMode is + selected. Override ``driver_station_connected()`` to run code once, when + the Driver Station first connects. + """ + + def __init__(self, *, auto_discover: bool = True): + """ + :param auto_discover: Automatically discover, register, and publish decorated + OpModes. If False, skip all automatic registration, + including already-imported decorated classes, and + automatic publication. + """ + super().__init__() + if auto_discover: + _opmode.discover_and_register(self) + + def add_opmode( + self, + opmode_cls: type, + mode: RobotMode, + name: str, + group: Optional[str] = None, + description: Optional[str] = None, + text_color: Optional[Color] = None, + background_color: Optional[Color] = None, + ) -> None: + """ + Adds an operating mode option. It's necessary to call publish_opmodes() to + make the added modes visible to the driver station. + + The text_color and background_color parameters are optional, but setting + only one has no effect (if only one is provided, it will be ignored). + + :param opmode_cls: opmode class; must be a public, non-abstract subclass of OpMode + with a constructor that either takes no arguments or accepts a + single argument of this class's type (the latter is preferred). + :param mode: robot mode + :param name: name of the operating mode + :param group: group of the operating mode + :param description: description of the operating mode + :param text_color: text color + :param background_color: background color + """ + + _opmode.register_opmode( + self, + opmode_cls, + mode, + name, + group, + description, + text_color, + background_color, + add_factory=self.add_opmode_factory, + )