probe: Support activate/deactivate scripts on each probe

Allow a set of g-code scripts to be run on each probe invocation.
This may be useful for probes that need to be setup before they are
useful (eg, with servo actuated probes).

Signed-off-by: Kevin O'Connor <kevin@koconnor.net>

3 files changed, 34 insertions, 0 deletions

diff --git a/klippy/extras/probe.py b/klippy/extras/probe.py
index aeebf725..b1e89944 100644
--- a/klippy/extras/probe.py
+++ b/klippy/extras/probe.py
@@ -16,6 +16,9 @@ class PrinterProbe:
         mcu = pin_params['chip']
         mcu.add_config_object(self)
         self.mcu_probe = mcu.setup_pin(pin_params)
+        if (config.get('activate_gcode', None) is not None or
+            config.get('deactivate_gcode', None) is not None):
+            self.mcu_probe = ProbeEndstopWrapper(config, self.mcu_probe)
         self.gcode = self.printer.lookup_object('gcode')
         self.gcode.register_command(
             'PROBE', self.cmd_PROBE, desc=self.cmd_PROBE_help)
@@ -49,6 +52,29 @@ class PrinterProbe:
         self.gcode.respond_info(
             "probe: %s" % (["open", "TRIGGERED"][not not res],))
 
+# Endstop wrapper that enables running g-code scripts on setup
+class ProbeEndstopWrapper:
+    def __init__(self, config, mcu_endstop):
+        self.mcu_endstop = mcu_endstop
+        self.gcode = config.get_printer().lookup_object('gcode')
+        self.activate_gcode = config.get('activate_gcode', "")
+        self.deactivate_gcode = config.get('deactivate_gcode', "")
+        # Wrappers
+        self.get_mcu = self.mcu_endstop.get_mcu
+        self.add_stepper = self.mcu_endstop.add_stepper
+        self.get_steppers = self.mcu_endstop.get_steppers
+        self.home_start = self.mcu_endstop.home_start
+        self.home_wait = self.mcu_endstop.home_wait
+        self.query_endstop = self.mcu_endstop.query_endstop
+        self.query_endstop_wait = self.mcu_endstop.query_endstop_wait
+        self.TimeoutError = self.mcu_endstop.TimeoutError
+    def home_prepare(self):
+        self.gcode.run_script(self.activate_gcode)
+        self.mcu_endstop.home_prepare()
+    def home_finalize(self):
+        self.gcode.run_script(self.deactivate_gcode)
+        self.mcu_endstop.home_finalize()
+
 # Helper code that can probe a series of points and report the
 # position at each point.
 class ProbePointsHelper:
diff --git a/klippy/homing.py b/klippy/homing.py
index 276c31d9..41f1ece5 100644
--- a/klippy/homing.py
+++ b/klippy/homing.py
@@ -42,6 +42,8 @@ class Homing:
         return dist_ticks / ticks_per_step
     def homing_move(self, movepos, endstops, speed, probe_pos=False):
         # Start endstop checking
+        for mcu_endstop, name in endstops:
+            mcu_endstop.home_prepare()
         print_time = self.toolhead.get_last_move_time()
         for mcu_endstop, name in endstops:
             min_step_dist = min([s.get_step_dist()
@@ -70,6 +72,8 @@ class Homing:
                 list(self.toolhead.get_kinematics().get_position()) + [None])
         else:
             self.toolhead.set_position(movepos)
+        for mcu_endstop, name in endstops:
+            mcu_endstop.home_finalize()
         if error is not None:
             raise EndstopError(error)
     def home(self, forcepos, movepos, endstops, speed, second_home=False):
diff --git a/klippy/mcu.py b/klippy/mcu.py
index 40fcbad7..06da3ea6 100644
--- a/klippy/mcu.py
+++ b/klippy/mcu.py
@@ -168,6 +168,8 @@ class MCU_endstop:
         self._query_cmd = self._mcu.lookup_command("end_stop_query oid=%c")
         self._mcu.register_msg(self._handle_end_stop_state, "end_stop_state"
                                , self._oid)
+    def home_prepare(self):
+        pass
     def home_start(self, print_time, sample_time, sample_count, rest_time):
         clock = self._mcu.print_time_to_clock(print_time)
         rest_ticks = int(rest_time * self._mcu.get_adjusted_freq())
@@ -184,6 +186,8 @@ class MCU_endstop:
         eventtime = self._mcu.monotonic()
         while self._check_busy(eventtime, home_end_time):
             eventtime = self._mcu.pause(eventtime + 0.1)
+    def home_finalize(self):
+        pass
     def _handle_end_stop_state(self, params):
         logging.debug("end_stop_state %s", params)
         self._last_state = params