From f72f94e299e33c77e6c1a4a75cf632fc8fc812e0 Mon Sep 17 00:00:00 2001
From: Kevin O'Connor <kevin@koconnor.net>
Date: Sat, 25 May 2024 20:52:18 -0400
Subject: [PATCH] probe: Move add_steppers() logic to HomingViaProbeHelper
 class

Perform the initial add_steppers() configuration in a single location.

Signed-off-by: Kevin O'Connor <kevin@koconnor.net>
---
 klippy/extras/bltouch.py            |  7 -------
 klippy/extras/probe.py              | 14 +++++++-------
 klippy/extras/probe_eddy_current.py |  7 -------
 3 files changed, 7 insertions(+), 21 deletions(-)

diff --git a/klippy/extras/bltouch.py b/klippy/extras/bltouch.py
index 482578c7d..081b92714 100644
--- a/klippy/extras/bltouch.py
+++ b/klippy/extras/bltouch.py
@@ -28,8 +28,6 @@ class BLTouchEndstopWrapper:
         self.printer = config.get_printer()
         self.printer.register_event_handler("klippy:connect",
                                             self.handle_connect)
-        self.printer.register_event_handler('klippy:mcu_identify',
-                                            self.handle_mcu_identify)
         self.position_endstop = config.getfloat('z_offset', minval=0.)
         self.stow_on_each_sample = config.getboolean('stow_on_each_sample',
                                                      True)
@@ -70,11 +68,6 @@ class BLTouchEndstopWrapper:
                                     desc=self.cmd_BLTOUCH_STORE_help)
         # multi probes state
         self.multi = 'OFF'
-    def handle_mcu_identify(self):
-        kin = self.printer.lookup_object('toolhead').get_kinematics()
-        for stepper in kin.get_steppers():
-            if stepper.is_active_axis('z'):
-                self.add_stepper(stepper)
     def handle_connect(self):
         self.sync_mcu_print_time()
         self.next_cmd_time += 0.200
diff --git a/klippy/extras/probe.py b/klippy/extras/probe.py
index 6f35f5e17..bb09217ac 100644
--- a/klippy/extras/probe.py
+++ b/klippy/extras/probe.py
@@ -179,6 +179,8 @@ class HomingViaProbeHelper:
         # Register z_virtual_endstop pin
         self.printer.lookup_object('pins').register_chip('probe', self)
         # Register event handlers
+        self.printer.register_event_handler('klippy:mcu_identify',
+                                            self._handle_mcu_identify)
         self.printer.register_event_handler("homing:homing_move_begin",
                                             self._handle_homing_move_begin)
         self.printer.register_event_handler("homing:homing_move_end",
@@ -189,6 +191,11 @@ class HomingViaProbeHelper:
                                             self._handle_home_rails_end)
         self.printer.register_event_handler("gcode:command_error",
                                             self._handle_command_error)
+    def _handle_mcu_identify(self):
+        kin = self.printer.lookup_object('toolhead').get_kinematics()
+        for stepper in kin.get_steppers():
+            if stepper.is_active_axis('z'):
+                self.mcu_probe.add_stepper(stepper)
     def _handle_homing_move_begin(self, hmove):
         if self.mcu_probe in hmove.get_mcu_endstops():
             self.mcu_probe.probe_prepare(hmove)
@@ -481,8 +488,6 @@ class ProbeEndstopWrapper:
         # Create an "endstop" object to handle the probe pin
         ppins = self.printer.lookup_object('pins')
         self.mcu_endstop = ppins.setup_pin('endstop', config.get('pin'))
-        self.printer.register_event_handler('klippy:mcu_identify',
-                                            self._handle_mcu_identify)
         # Wrappers
         self.get_mcu = self.mcu_endstop.get_mcu
         self.add_stepper = self.mcu_endstop.add_stepper
@@ -492,11 +497,6 @@ class ProbeEndstopWrapper:
         self.query_endstop = self.mcu_endstop.query_endstop
         # multi probes state
         self.multi = 'OFF'
-    def _handle_mcu_identify(self):
-        kin = self.printer.lookup_object('toolhead').get_kinematics()
-        for stepper in kin.get_steppers():
-            if stepper.is_active_axis('z'):
-                self.add_stepper(stepper)
     def _raise_probe(self):
         toolhead = self.printer.lookup_object('toolhead')
         start_pos = toolhead.get_position()
diff --git a/klippy/extras/probe_eddy_current.py b/klippy/extras/probe_eddy_current.py
index 3f4a5e206..d2309617f 100644
--- a/klippy/extras/probe_eddy_current.py
+++ b/klippy/extras/probe_eddy_current.py
@@ -196,13 +196,6 @@ class EddyEndstopWrapper:
         self._samples = []
         self._is_sampling = self._start_from_home = self._need_stop = False
         self._trigger_time = 0.
-        self._printer.register_event_handler('klippy:mcu_identify',
-                                             self._handle_mcu_identify)
-    def _handle_mcu_identify(self):
-        kin = self._printer.lookup_object('toolhead').get_kinematics()
-        for stepper in kin.get_steppers():
-            if stepper.is_active_axis('z'):
-                self.add_stepper(stepper)
     # Measurement gathering
     def _start_measurements(self, is_home=False):
         self._need_stop = False