tmc: Query the stepper step_distance from the stepper object

Don't obtain the step_distance from the config, instead obtain
step_distance from the instantiated stepper object.

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

2 files changed, 28 insertions, 17 deletions

diff --git a/klippy/extras/force_move.py b/klippy/extras/force_move.py
index d051622f..6c4a6025 100644
--- a/klippy/extras/force_move.py
+++ b/klippy/extras/force_move.py
@@ -53,6 +53,10 @@ class ForceMove:
     def register_stepper(self, stepper):
         name = stepper.get_name()
         self.steppers[name] = stepper
+    def lookup_stepper(self, name):
+        if name not in self.steppers:
+            raise self.printer.config_error("Unknown stepper %s" % (name,))
+        return self.steppers[name]
     def force_enable(self, stepper):
         toolhead = self.printer.lookup_object('toolhead')
         print_time = toolhead.get_last_move_time()
diff --git a/klippy/extras/tmc.py b/klippy/extras/tmc.py
index a9d99cc0..5f0f5f7e 100644
--- a/klippy/extras/tmc.py
+++ b/klippy/extras/tmc.py
@@ -274,21 +274,28 @@ class TMCMicrostepHelper:
         return (1023 - mscnt) >> self.fields.get_field("MRES")
 
 # Helper to configure "stealthchop" mode
-def TMCStealthchopHelper(config, mcu_tmc, tmc_freq):
-    fields = mcu_tmc.get_fields()
-    en_pwm_mode = False
-    velocity = config.getfloat('stealthchop_threshold', 0., minval=0.)
-    if velocity:
-        stepper_name = " ".join(config.get_name().split()[1:])
-        stepper_config = config.getsection(stepper_name)
-        step_dist = stepper_config.getfloat('step_distance')
+class TMCStealthchopHelper:
+    def __init__(self, config, mcu_tmc, tmc_freq):
+        self.name = config.get_name()
+        self.fields = fields = mcu_tmc.get_fields()
+        self.tmc_freq = tmc_freq
+        printer = config.get_printer()
+        self.force_move = printer.try_load_module(config, "force_move")
+        en_pwm_mode = False
+        self.velocity = config.getfloat('stealthchop_threshold', 0., minval=0.)
+        if self.velocity:
+            printer.register_event_handler("connect", self.handle_connect)
+            en_pwm_mode = True
+        reg = fields.lookup_register("en_pwm_mode", None)
+        if reg is not None:
+            fields.set_field("en_pwm_mode", en_pwm_mode)
+        else:
+            # TMC2208 uses en_spreadCycle
+            fields.set_field("en_spreadCycle", not en_pwm_mode)
+    def handle_connect(self):
+        stepper_name = " ".join(self.name.split()[1:])
+        stepper = self.force_move.lookup_stepper(stepper_name)
+        step_dist = stepper.get_step_dist()
         step_dist_256 = step_dist / (1 << fields.get_field("MRES"))
-        threshold = int(tmc_freq * step_dist_256 / velocity + .5)
-        fields.set_field("TPWMTHRS", max(0, min(0xfffff, threshold)))
-        en_pwm_mode = True
-    reg = fields.lookup_register("en_pwm_mode", None)
-    if reg is not None:
-        fields.set_field("en_pwm_mode", en_pwm_mode)
-    else:
-        # TMC2208 uses en_spreadCycle
-        fields.set_field("en_spreadCycle", not en_pwm_mode)
+        threshold = int(self.tmc_freq * step_dist_256 / self.velocity + .5)
+        self.fields.set_field("TPWMTHRS", max(0, min(0xfffff, threshold)))