[py] Fix opmodes (#8498)

Co-authored-by: David Vo <auscompgeek@users.noreply.github.com>

hal/robotpy_pybind_build_info.bzl

@@ -333,6 +333,7 @@ def wpihal_extension(srcs = [], header_to_dat_deps = [], extra_hdrs = [], includ
             header_file = "$(execpath :robotpy-native-wpihal.copy_headers)/wpi/hal/DriverStationTypes.h",
             tmpl_class_names = [],
             trampolines = [
+                ("HAL_ControlWord", "__HAL_ControlWord.hpp"),
                 ("HAL_JoystickAxes", "__HAL_JoystickAxes.hpp"),
                 ("HAL_JoystickPOVs", "__HAL_JoystickPOVs.hpp"),
                 ("HAL_JoystickButtons", "__HAL_JoystickButtons.hpp"),

hal/src/main/python/semiwrap/DriverStation.yml

@@ -42,6 +42,4 @@ functions:
   HAL_SetOpModeOptions:
   HAL_ObserveUserProgram:
   GetControlWord:
-    ignore: true
   GetUncachedControlWord:
-    ignore: true

hal/src/main/python/semiwrap/DriverStationTypes.yml

@@ -12,11 +12,13 @@ enums:
   HAL_MatchType:
     value_prefix: HAL_kMatchType
   HAL_RobotMode:
-    ignore: true
+    rename: _RobotMode
   RobotMode:
 classes:
   HAL_ControlWord:
-    ignore: true
+    rename: _ControlWord
+    attributes:
+      value:
   HAL_JoystickAxes:
     attributes:
       axes:

hal/src/main/python/semiwrap/simulation/DriverStationData.yml

@@ -127,6 +127,4 @@ functions:
     ignore: true
   HALSIM_CancelOpModeOptionsCallback:
   HALSIM_GetOpModeOptions:
-    ignore: true
   HALSIM_FreeOpModeOptionsArray:
-    ignore: true

hal/src/main/python/semiwrap/simulation/MockHooks.yml

@@ -11,9 +11,7 @@ functions:
   HALSIM_SetProgramStarted:
   HALSIM_GetProgramStarted:
   HALSIM_SetProgramState:
-    ignore: true
   HALSIM_GetProgramState:
-    ignore: true
   SetProgramState:
   GetProgramState:
   HALSIM_RestartTiming:

shared/bazel/rules/robotpy/pytest_util.bzl

@@ -1,13 +1,13 @@
 load("@rules_python_pytest//python_pytest:defs.bzl", "py_pytest_test")
 load("//shared/bazel/rules/robotpy:compatibility_select.bzl", "robotpy_compatibility_select")
 
-def robotpy_py_test(name, srcs, **kwargs):
+def robotpy_py_test(name, srcs, tags = [], **kwargs):
     py_pytest_test(
         name = name,
         size = "small",
         srcs = srcs,
         target_compatible_with = robotpy_compatibility_select(),
-        tags = [
+        tags = tags + [
             "no-asan",
             "no-tsan",
             "robotpy",

wpilibc/BUILD.bazel

@@ -291,6 +291,7 @@ define_pybind_library(
 robotpy_py_test(
     "python_tests",
     srcs = glob(["src/test/python/**/*.py"]),
+    tags = ["exclusive"],
     deps = [
         ":robotpy-wpilib",
         requirement("pytest"),

wpilibc/robotpy_pybind_build_info.bzl

@@ -1513,6 +1513,7 @@ def wpilib_simulation_extension(srcs = [], header_to_dat_deps = [], extra_hdrs =
             header_file = "$(execpath :robotpy-native-wpilib.copy_headers)/wpi/simulation/DriverStationSim.hpp",
             tmpl_class_names = [],
             trampolines = [
+                ("wpi::sim::OpModeOptions", "wpi__sim__OpModeOptions.hpp"),
                 ("wpi::sim::DriverStationSim", "wpi__sim__DriverStationSim.hpp"),
             ],
         ),

wpilibc/src/main/python/semiwrap/OpModeRobot.yml

@@ -19,8 +19,3 @@ classes:
       ClearOpModes:
   wpi::OpModeRobot:
     ignore: true
-    methods:
-      AddOpMode:
-        overloads:
-          RobotMode, std::string_view, std::string_view, std::string_view, const wpi::util::Color&, const wpi::util::Color&:
-          RobotMode, std::string_view, std::string_view, std::string_view:

wpilibc/src/main/python/semiwrap/simulation/DriverStationSim.yml

@@ -1,6 +1,27 @@
 classes:
   wpi::sim::OpModeOptions:
-    ignore: true
+    ignored_bases:
+    - std::span<HAL_OpModeOption>
+    force_no_trampoline: true
+    methods:
+      OpModeOptions:
+        overloads:
+          "":
+            ignore: true
+          HAL_OpModeOption*, int32_t:
+            ignore: true
+    inline_code: |
+      .def("__len__", [](const OpModeOptions &self) { return self.size(); })
+      .def("__getitem__", [](const OpModeOptions &self, int index) {
+        if (index >= static_cast<int>(self.size())) {
+          throw std::out_of_range("OpModeOptions index out of range");
+        }
+        return self[index];
+      })
+      .def("__iter__", [](OpModeOptions &self) {
+        return py::make_iterator(self.begin(), self.end());
+      }, py::keep_alive<0,1>());
+
   wpi::sim::DriverStationSim:
     force_type_casters:
     - std::function
@@ -54,5 +75,5 @@ classes:
       GetOpMode:
       SetOpMode:
       RegisterOpModeOptionsCallback:
-      GetOpModeOptions:
         ignore: true
+      GetOpModeOptions:

wpilibc/src/main/python/wpilib/__init__.py

@@ -47,7 +47,6 @@ from ._wpilib import (
     OnboardIMU,
     OpMode,
     OpModeRobotBase,
-    PeriodicOpMode,
     PS4Controller,
     PS5Controller,
     PWM,
@@ -58,6 +57,7 @@ from ._wpilib import (
     PWMTalonSRX,
     PWMVenom,
     PWMVictorSPX,
+    PeriodicOpMode,
     PneumaticHub,
     PneumaticsBase,
     PneumaticsControlModule,
@@ -139,7 +139,6 @@ __all__ = [
     "OnboardIMU",
     "OpMode",
     "OpModeRobotBase",
-    "PeriodicOpMode",
     "PS4Controller",
     "PS5Controller",
     "PWM",
@@ -150,6 +149,7 @@ __all__ = [
     "PWMTalonSRX",
     "PWMVenom",
     "PWMVictorSPX",
+    "PeriodicOpMode",
     "PneumaticHub",
     "PneumaticsBase",
     "PneumaticsControlModule",

wpilibc/src/main/python/wpilib/simulation/__init__.py

@@ -32,6 +32,7 @@ from ._simulation import (
     LinearSystemSim_2_1_2,
     LinearSystemSim_2_2_1,
     LinearSystemSim_2_2_2,
+    OpModeOptions,
     PS4ControllerSim,
     PS5ControllerSim,
     PWMMotorControllerSim,
@@ -48,11 +49,13 @@ from ._simulation import (
     StadiaControllerSim,
     XboxControllerSim,
     getProgramStarted,
+    getProgramState,
     isTimingPaused,
     pauseTiming,
     restartTiming,
     resumeTiming,
     setProgramStarted,
+    setProgramState,
     setRuntimeType,
     stepTiming,
     stepTimingAsync,
@@ -87,6 +90,7 @@ __all__ = [
     "LinearSystemSim_2_1_2",
     "LinearSystemSim_2_2_1",
     "LinearSystemSim_2_2_2",
+    "OpModeOptions",
     "PS4ControllerSim",
     "PS5ControllerSim",
     "PWMMotorControllerSim",
@@ -103,11 +107,13 @@ __all__ = [
     "StadiaControllerSim",
     "XboxControllerSim",
     "getProgramStarted",
+    "getProgramState",
     "isTimingPaused",
     "pauseTiming",
     "restartTiming",
     "resumeTiming",
     "setProgramStarted",
+    "setProgramState",
     "setRuntimeType",
     "stepTiming",
     "stepTimingAsync",

wpilibc/src/test/python/test_opmode_robot.py

@@ -0,0 +1,156 @@
+import pytest
+import threading
+from wpilib import simulation as wsim
+from wpimath.units import seconds
+from wpilib.opmoderobot import OpModeRobot
+from wpilib import OpMode
+from hal._wpiHal import RobotMode
+from wpiutil import Color
+
+
+class MockOpMode(OpMode):
+    def __init__(self):
+        super().__init__()
+        self.disabled_periodic_count = 0
+        self.op_mode_run_count = 0
+        self.op_mode_stop_count = 0
+
+    def disabled_periodic(self):
+        self.disabled_periodic_count += 1
+
+    def op_mode_run(self, op_mode_id: int):
+        self.op_mode_run_count += 1
+
+    def op_mode_stop(self):
+        self.op_mode_stop_count += 1
+
+
+class OneArgOpMode(OpMode):
+    def __init__(self, robot):
+        super().__init__()
+
+    def op_mode_run(self, op_mode_id: int):
+        pass
+
+    def op_mode_stop(self):
+        pass
+
+
+class MockRobot(OpModeRobot):
+    def __init__(self):
+        super().__init__()
+        self.driver_station_connected_count = 0
+        self.none_periodic_count = 0
+
+    def driverStationConnected(self):
+        self.driver_station_connected_count += 1
+
+    def nonePeriodic(self):
+        self.none_periodic_count += 1
+
+
+@pytest.fixture(autouse=True)
+def sim_timing_setup():
+    wsim.pauseTiming()
+    wsim.setProgramStarted(False)
+    yield
+    wsim.resumeTiming()
+
+
+def test_add_op_mode():
+    class MyMockRobot(MockRobot):
+        def __init__(self):
+            super().__init__()
+            self.addOpMode(
+                MockOpMode,
+                RobotMode.AUTONOMOUS,
+                "NoArgOpMode-Auto",
+                "Group",
+                "Description",
+                Color.kWhite,
+                Color.kBlack,
+            )
+            self.addOpMode(
+                OneArgOpMode,
+                RobotMode.TEST,
+                "OneArgOpMode-Test",
+                "Group",
+                "Description",
+                Color.kWhite,
+                Color.kBlack,
+            )
+            self.addOpMode(MockOpMode, RobotMode.TELEOPERATED, "NoArgOpMode")
+            self.addOpMode(OneArgOpMode, RobotMode.TELEOPERATED, "OneArgOpMode")
+            self.publishOpModes()
+
+    robot = MyMockRobot()
+    options = wsim.DriverStationSim.getOpModeOptions()
+
+    assert len(options) == 4
+
+    opt_map = {opt.name: opt for opt in options}
+
+    auto_opt = opt_map["NoArgOpMode-Auto"]
+    assert auto_opt.group == "Group"
+    assert auto_opt.description == "Description"
+    assert auto_opt.textColor == 0xFFFFFF
+    assert auto_opt.backgroundColor == 0x000000
+
+    tele_opt = opt_map["NoArgOpMode"]
+    assert tele_opt.group == ""
+    assert tele_opt.description == ""
+    assert tele_opt.textColor == -1
+    assert tele_opt.backgroundColor == -1
+
+
+def test_clear_op_modes():
+    class MyMockRobot(MockRobot):
+        def __init__(self):
+            super().__init__()
+            self.addOpMode(MockOpMode, RobotMode.TELEOPERATED, "NoArgOpMode")
+            self.publishOpModes()
+
+    robot = MyMockRobot()
+    robot.clearOpModes()
+
+    options = wsim.DriverStationSim.getOpModeOptions()
+    assert len(options) == 0
+
+
+def test_remove_op_mode():
+    class MyMockRobot(MockRobot):
+        def __init__(self):
+            super().__init__()
+            self.addOpMode(MockOpMode, RobotMode.TELEOPERATED, "NoArgOpMode")
+            self.addOpMode(OneArgOpMode, RobotMode.TELEOPERATED, "OneArgOpMode")
+            self.publishOpModes()
+
+    robot = MyMockRobot()
+    robot.removeOpMode(RobotMode.TELEOPERATED, "NoArgOpMode")
+    robot.publishOpModes()
+
+    options = wsim.DriverStationSim.getOpModeOptions()
+    assert len(options) == 1
+    assert options[0].name == "OneArgOpMode"
+
+
+def test_none_periodic():
+    class MyMockRobot(MockRobot):
+        def __init__(self):
+            super().__init__()
+            self.addOpMode(MockOpMode, RobotMode.TELEOPERATED, "NoArgOpMode")
+            self.publishOpModes()
+
+    robot = MyMockRobot()
+
+    robot_thread = threading.Thread(target=robot.startCompetition)
+    robot_thread.start()
+
+    wsim.waitForProgramStart()
+
+    wsim.stepTiming(0.110)
+
+    assert robot.none_periodic_count == 2
+
+    robot.endCompetition()
+    robot_thread.join()