Fix bug where autonomous mode would repeat if ended with a timed state

- Also guarantee that done() will always be called when autonomous mode ends, even if engage() is continually called
- Fixes #108

Author: virtuald

magicbot/state_machine.py

@@ -479,6 +479,9 @@ def next_state_now(self, name):
     def done(self):
         '''Call this function to end execution of the state machine.
         
+        This function will always be called when a state machine ends. Even if
+        the engage function is called repeatedly, done() will be called.
+        
         .. note:: If you wish to do something each time execution ceases,
                   override this function (but be sure to call
                   ``super().done()``!)
@@ -512,6 +515,7 @@ def execute(self):
         # tm is the number of seconds that the state machine has been executing
         tm = now - self.__start
         state = self.__state
+        done_called = False
 
         # we adjust this so that if we have states chained together,
         # then the total time it runs is the amount of time of the
@@ -526,6 +530,10 @@ def execute(self):
             if state.next_state is None:
                 # If the state expires and it's the last state, if the machine
                 # is still engaged then it should cycle back to the beginning
+                # ... but we should call done() first
+                done_called = True
+                self.done()
+                
                 if self.__should_engage:
                     self.next_state(self.__first)
                     state = self.__state
@@ -561,7 +569,7 @@ def execute(self):
 
             # execute the state function, passing it the arguments
             state.run(self, tm, tm - state.start_time, initial_call)
-        else:
+        elif not done_called:
             # or clear the state
             self.done()
 
@@ -601,3 +609,8 @@ def on_iteration(self, tm):
             self.engage()
             self.execute()
             self.__engaged = self.is_executing
+    
+    def done(self):
+        super().done()
+        self._StateMachine__should_engage = False
+        self.__engaged = False

tests/test_magicbot_sm.py

@@ -259,6 +259,43 @@ def something(self):
         sm.on_iteration(None)
         assert not sm.is_executing
 
+    assert sm.i == 6
+
+def test_autonomous_sm_end_timed_state(wpitime):
+    class _TM(AutonomousStateMachine):
+        
+        i = 0
+        j = 0
+        VERBOSE_LOGGING = False
+        
+        @state(first=True)
+        def something(self):
+            self.i += 1
+            if self.i == 3:
+                self.next_state('timed')
+        
+        @timed_state(duration=1)
+        def timed(self):
+            self.j += 1
+        
+    sm = _TM()
+    setup_tunables(sm, 'cname')
+    
+    sm.on_enable()
+    
+    for _ in range(5):
+        wpitime.now += 0.7
+        sm.on_iteration(None)
+        assert sm.is_executing
+    
+    for _ in range(5):
+        wpitime.now += 0.7
+        sm.on_iteration(None)
+        assert not sm.is_executing
+    
+    assert sm.i == 3
+    assert sm.j == 2
+
 def test_next_fn():
     class _TM(StateMachine):
         @state(first=True)
@@ -528,6 +565,10 @@ def a(self):
         @timed_state(duration=.01)
         def b(self):
             self.executed.append('b')
+            
+        def done(self):
+            super().done()
+            self.executed.append('d')
 
     sm = _SM()
     setup_tunables(sm, 'cname')
@@ -547,4 +588,4 @@ def b(self):
 
         wpitime.now += 0.02
 
-    assert sm.executed == ['a', 'b', 'a', 'b', 'a', 'b', 'a', 'b']
+    assert sm.executed == ['a', 'b', 'd', 'a', 'b', 'd', 'a', 'b', 'd', 'a', 'b']