Merge pull request #201 from robotpy/fix-magicbot-ds

Update magicbot for new DriverStation threading model

Author: virtuald

magicbot/magicrobot.py

@@ -377,6 +377,7 @@ def startCompetition(self) -> None:
 
     def endCompetition(self) -> None:
         self.__done = True
+        self._automodes.endCompetition()
 
     def autonomous(self) -> None:
         """
@@ -434,9 +435,17 @@ def _disabled(self) -> None:
             self.onException(forceReport=True)
         watchdog.addEpoch("disabledInit()")
 
+        refreshData = wpilib.DriverStation.refreshData
+        DSControlWord = wpilib.DSControlWord
+
         with NotifierDelay(self.control_loop_wait_time) as delay:
-            while not self.__done and self.isDisabled():
-                if ds_attached != self.__is_ds_attached():
+            while not self.__done:
+                refreshData()
+                cw = DSControlWord()
+                if cw.isEnabled():
+                    break
+
+                if ds_attached != cw.isDSAttached():
                     ds_attached = not ds_attached
                     self.__nt_put_is_ds_attached(ds_attached)
 
@@ -484,9 +493,15 @@ def _operatorControl(self) -> None:
         watchdog.addEpoch("teleopInit()")
 
         observe = hal.observeUserProgramTeleop
+        refreshData = wpilib.DriverStation.refreshData
+        isTeleopEnabled = wpilib.DriverStation.isTeleopEnabled
 
         with NotifierDelay(self.control_loop_wait_time) as delay:
-            while not self.__done and self.isTeleopEnabled():
+            while not self.__done:
+                refreshData()
+                if not isTeleopEnabled():
+                    break
+
                 observe()
                 try:
                     self.teleopPeriodic()
@@ -526,8 +541,16 @@ def _test(self) -> None:
             self.onException(forceReport=True)
         watchdog.addEpoch("testInit()")
 
+        refreshData = wpilib.DriverStation.refreshData
+        DSControlWord = wpilib.DSControlWord
+
         with NotifierDelay(self.control_loop_wait_time) as delay:
-            while not self.__done and self.isTest() and self.isEnabled():
+            while not self.__done:
+                refreshData()
+                cw = DSControlWord()
+                if not (cw.isTest() and cw.isEnabled()):
+                    break
+
                 hal.observeUserProgramTest()
                 try:
                     self.testPeriodic()

robotpy_ext/autonomous/selector.py

@@ -80,6 +80,7 @@ def __init__(self, autonomous_pkgname, *args, **kwargs):
 
         self.modes = {}
         self.active_mode = None
+        self.robot_exit = False
 
         logger.info("Begin initializing autonomous mode switcher")
 
@@ -210,6 +211,10 @@ def __init__(self, autonomous_pkgname, *args, **kwargs):
 
         logger.info("Autonomous switcher initialized")
 
+    def endCompetition(self):
+        """Call this function when your robot's endCompetition function is called"""
+        self.robot_exit = True
+
     def run(
         self,
         control_loop_wait_time: float = 0.020,
@@ -272,10 +277,15 @@ def watchdog_check_expired():
         #
 
         observe = hal.observeUserProgramAutonomous
+        refreshData = wpilib.DriverStation.refreshData
         isAutonomousEnabled = wpilib.DriverStation.isAutonomousEnabled
 
         with NotifierDelay(control_loop_wait_time) as delay:
-            while isAutonomousEnabled():
+            while not self.robot_exit:
+                refreshData()
+                if not isAutonomousEnabled():
+                    break
+
                 observe()
                 try:
                     self._on_iteration(timer.get())