Add restart grace period to multiball

mpf/config_spec.yaml

@@ -1141,6 +1141,7 @@ multiballs:
     shoot_again: single|template_ms|10s
     hurry_up_time: single|template_ms|0
     grace_period: single|template_ms|0
+    restart_grace_period: single|template_ms|0
     ball_locks: list|machine(ball_devices)|None
     enable_events: event_handler|event_handler:ms|None
     disable_events: event_handler|event_handler:ms|None

mpf/devices/multiball.py

@@ -237,15 +237,35 @@ def _ball_drain_count_balls(self, balls, **kwargs):
         desc: The multiball called (name) has lost a ball after ball save expired.
         '''
 
-        if not self.machine.game or self.machine.game.balls_in_play - balls < 1:
-            self.balls_added_live = 0
-            self.balls_live_target = 0
-            self.machine.events.remove_handler(self._ball_drain_count_balls)
-            self.machine.events.post("multiball_{}_ended".format(self.name))
-            '''event: multiball_(name)_ended
-            desc: The multiball called (name) has just ended.
-            '''
-            self.debug_log("Ball drained. MB ended.")
+        restart_grace_period_ms = self.config['restart_grace_period'].evaluate([])
+
+        if not self.machine.game or (self.machine.game.balls_in_play - balls < 1 and restart_grace_period_ms==0):
+            self._multiball_end()
+        elif restart_grace_period_ms>0:
+            if self.machine.game.balls_in_play == 0:
+                self.delay.run_now(f"{self.name}_restart_grace_period")
+            elif self.machine.game.balls_in_play - balls < 1:
+                self.machine.events.post(f"multiball_{self.name}_will_end", duration=restart_grace_period_ms)
+                '''event: multiball_(name)_will_end
+                desc: The multiball called (name) does not have multiple balls active, enabling grace period.
+                args:
+                    duration: duration of the restart grace period
+                '''
+
+                self.delay.add(name=f"{self.name}_restart_grace_period",
+                               ms=restart_grace_period_ms,
+                               callback=self._multiball_end)
+
+    def _multiball_end(self):
+        self.balls_added_live = 0
+        self.balls_live_target = 0
+        self.machine.events.remove_handler(self._ball_drain_count_balls)
+        self.machine.events.post("multiball_{}_ended".format(self.name))
+        '''event: multiball_(name)_ended
+        desc: The multiball called (name) has just ended.
+        '''
+        self.debug_log("Ball drained. MB ended.")
+
 
     @event_handler(5)
     def event_stop(self, **kwargs):
@@ -284,6 +304,15 @@ def event_add_a_ball(self, **kwargs):
 
     def add_a_ball(self):
         """Add a ball if multiball has started."""
+
+        # remove if it exists as multiball now has multiple balls
+        if self.delay.check(f"{self.name}_restart_grace_period"):
+            self.delay.remove(f"{self.name}_restart_grace_period")
+            self.machine.events.post(f"multiball_{self.name}_will_resume")
+            '''event: multiball_(name)_will_resume
+            desc: Multiball (name) has been resumed during end of multiball grace period.
+            '''
+
         if self.balls_live_target > 0:
             self.debug_log("Adding a ball.")
             self.balls_live_target += 1

mpf/tests/test_MultiBall.py

@@ -1228,3 +1228,52 @@ def testAddABallSaverDuringShootAgain(self):
         self.assertAvailableBallsOnPlayfield(3)
         self.assertEventCalled("ball_save_mb_add_a_ball_timers_timer_start", 1)
         self.assertEventNotCalled("ball_save_mb_add_a_ball_timers_add_a_ball_timer_start")
+
+    def testEndMultiballGracePeriodNoRestart(self):
+        self.fill_troughs()
+        self.start_game()
+        self.assertAvailableBallsOnPlayfield(1)
+        self.mock_event("multiball_mb_grace_period_restart_will_end")
+        self.mock_event("multiball_mb_grace_period_restart_ended")
+
+        # start mb, default ball save, 5s restart grace period
+        self.post_event("mb_grace_period_restart_start")
+        self.advance_time_and_run(10)
+        self.assertAvailableBallsOnPlayfield(2)
+
+        self.drain_one_ball()
+        self.advance_time_and_run(5)
+
+        self.assertEventCalled("multiball_mb_grace_period_restart_will_end", 1)
+        self.assertEventCalled("multiball_mb_grace_period_restart_ended", 0)
+
+        self.advance_time_and_run(5)
+        self.assertEventCalled("multiball_mb_grace_period_restart_ended", 1)
+
+
+    def testEndMultiballGracePeriodAddABallRestart(self):
+        self.fill_troughs()
+        self.start_game()
+        self.assertAvailableBallsOnPlayfield(1)
+        self.mock_event("multiball_mb_grace_period_restart_will_end")
+        self.mock_event("multiball_mb_grace_period_restart_ended")
+        self.mock_event("multiball_mb_grace_period_restart_will_resume")
+        self.mock_event("ball_save_mb_grace_period_restart_timer_start")
+
+        # start mb, default ball save, 5s restart grace period
+        self.post_event("mb_grace_period_restart_start")
+        self.advance_time_and_run(10)
+        self.assertAvailableBallsOnPlayfield(2)
+
+        self.drain_one_ball()
+        self.advance_time_and_run(3)
+
+        self.assertEventCalled("multiball_mb_grace_period_restart_will_end", 1)
+        self.assertEventCalled("multiball_mb_grace_period_restart_ended", 0)
+
+        self.post_event("add_ball")
+        self.assertEventCalled("multiball_mb_grace_period_restart_ended", 0)
+        self.assertEventCalled("multiball_mb_grace_period_restart_will_resume", 1)
+        self.advance_time_and_run(5)
+        self.assertAvailableBallsOnPlayfield(2)
+        self.assertEventCalled("ball_save_mb_grace_period_restart_timer_start")