homing: Pass the HomingMove class to homing_move_begin/end events

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

1 files changed, 8 insertions, 8 deletions

diff --git a/klippy/extras/probe.py b/klippy/extras/probe.py
index 7b4ddcb7..e4f7a9ac 100644
--- a/klippy/extras/probe.py
+++ b/klippy/extras/probe.py
@@ -70,12 +70,12 @@ class PrinterProbe:
                                     desc=self.cmd_PROBE_CALIBRATE_help)
         self.gcode.register_command('PROBE_ACCURACY', self.cmd_PROBE_ACCURACY,
                                     desc=self.cmd_PROBE_ACCURACY_help)
-    def _handle_homing_move_begin(self, endstops):
-        if self.mcu_probe in endstops:
-            self.mcu_probe.probe_prepare()
-    def _handle_homing_move_end(self, endstops):
-        if self.mcu_probe in endstops:
-            self.mcu_probe.probe_finish()
+    def _handle_homing_move_begin(self, hmove):
+        if self.mcu_probe in hmove.get_mcu_endstops():
+            self.mcu_probe.probe_prepare(hmove)
+    def _handle_homing_move_end(self, hmove):
+        if self.mcu_probe in hmove.get_mcu_endstops():
+            self.mcu_probe.probe_finish(hmove)
     def _handle_home_rails_begin(self, homing_state, rails):
         endstops = [es for rail in rails for es, name in rail.get_endstops()]
         if self.mcu_probe in endstops:
@@ -296,14 +296,14 @@ class ProbeEndstopWrapper:
         pass
     def multi_probe_end(self):
         pass
-    def probe_prepare(self):
+    def probe_prepare(self, hmove):
         toolhead = self.printer.lookup_object('toolhead')
         start_pos = toolhead.get_position()
         self.activate_gcode.run_gcode_from_command()
         if toolhead.get_position()[:3] != start_pos[:3]:
             raise self.printer.command_error(
                 "Toolhead moved during probe activate_gcode script")
-    def probe_finish(self):
+    def probe_finish(self, hmove):
         toolhead = self.printer.lookup_object('toolhead')
         start_pos = toolhead.get_position()
         self.deactivate_gcode.run_gcode_from_command()