Add restart grace period to multiball

Allow multiballs to be restarted during grace period via add-a-ball.

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,33 @@ 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"multiball_{self.name}_grace_period_restart")
+            elif self.machine.game.balls_in_play - balls < 1:
+                self.machine.events.post("multiball_{}_will_end".format(self.name))
+                '''event: multiball_(name)_will_end
+                desc: The multiball called (name) does not have multiple balls active, enabling grace period.
+                '''
+
+                self.delay.add(name=f"multiball_{self.name}_grace_period_restart",
+                               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 +302,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"multiball_{self.name}_grace_period_restart"):
+            self.delay.remove(f"multiball_{self.name}_grace_period_restart")
+            self.machine.events.post("multiball_{}_will_resume".format(self.name))
+            '''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