tmc: Move SET_TMC_CURRENT command to TMCCommandHelper()

Refactor the tmc driver implementations so that there is a single
implementation of the SET_TMC_CURRENT command.

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

6 files changed, 37 insertions, 68 deletions

diff --git a/klippy/extras/tmc.py b/klippy/extras/tmc.py
index fbd9f584..fa2c0b57 100644
--- a/klippy/extras/tmc.py
+++ b/klippy/extras/tmc.py
@@ -81,11 +81,12 @@ class FieldHelper:
 ######################################################################
 
 class TMCCommandHelper:
-    def __init__(self, config, mcu_tmc):
+    def __init__(self, config, mcu_tmc, current_helper):
         self.printer = config.get_printer()
         self.stepper_name = ' '.join(config.get_name().split()[1:])
         self.name = config.get_name().split()[-1]
         self.mcu_tmc = mcu_tmc
+        self.current_helper = current_helper
         self.fields = mcu_tmc.get_fields()
         self.read_registers = self.read_translate = None
         self.toff = None
@@ -99,6 +100,9 @@ class TMCCommandHelper:
         gcode.register_mux_command("INIT_TMC", "STEPPER", self.name,
                                    self.cmd_INIT_TMC,
                                    desc=self.cmd_INIT_TMC_help)
+        gcode.register_mux_command("SET_TMC_CURRENT", "STEPPER", self.name,
+                                   self.cmd_SET_TMC_CURRENT,
+                                   desc=self.cmd_SET_TMC_CURRENT_help)
     def _init_registers(self, print_time=None):
         # Send registers
         for reg_name, val in self.fields.registers.items():
@@ -141,6 +145,28 @@ class TMCCommandHelper:
         reg_val = self.fields.set_field(field_name, value)
         print_time = self.printer.lookup_object('toolhead').get_last_move_time()
         self.mcu_tmc.set_register(reg_name, reg_val, print_time)
+    cmd_SET_TMC_CURRENT_help = "Set the current of a TMC driver"
+    def cmd_SET_TMC_CURRENT(self, gcmd):
+        ch = self.current_helper
+        prev_run_current, prev_hold_current, max_current = ch.get_current()
+        run_current = gcmd.get_float('CURRENT', None,
+                                     minval=0., maxval=max_current)
+        hold_current = gcmd.get_float('HOLDCURRENT', None,
+                                      above=0., maxval=max_current)
+        if run_current is None and hold_current is None:
+            # Query only
+            if prev_hold_current is None:
+                gcmd.respond_info("Run Current: %0.2fA" % (prev_run_current,))
+            else:
+                gcmd.respond_info("Run Current: %0.2fA Hold Current: %0.2fA"
+                                  % (prev_run_current, prev_hold_current))
+            return
+        if run_current is None:
+            run_current = prev_run_current
+        if hold_current is None:
+            hold_current = prev_hold_current
+        print_time = self.printer.lookup_object('toolhead').get_last_move_time()
+        ch.set_current(run_current, hold_current, print_time)
     # Stepper enable/disable via comms
     def _do_enable(self, print_time, is_enable):
         toff_val = 0
diff --git a/klippy/extras/tmc2130.py b/klippy/extras/tmc2130.py
index 0f503726..17467ae1 100644
--- a/klippy/extras/tmc2130.py
+++ b/klippy/extras/tmc2130.py
@@ -108,10 +108,6 @@ class TMCCurrentHelper:
         self.fields.set_field("vsense", vsense)
         self.fields.set_field("IHOLD", ihold)
         self.fields.set_field("IRUN", irun)
-        gcode = self.printer.lookup_object("gcode")
-        gcode.register_mux_command("SET_TMC_CURRENT", "STEPPER", self.name,
-                                   self.cmd_SET_TMC_CURRENT,
-                                   desc=self.cmd_SET_TMC_CURRENT_help)
     def _calc_current_bits(self, current, vsense):
         sense_resistor = self.sense_resistor + 0.020
         vref = 0.32
@@ -150,24 +146,6 @@ class TMCCurrentHelper:
         self.fields.set_field("IHOLD", ihold)
         val = self.fields.set_field("IRUN", irun)
         self.mcu_tmc.set_register("IHOLD_IRUN", val, print_time)
-    cmd_SET_TMC_CURRENT_help = "Set the current of a TMC driver"
-    def cmd_SET_TMC_CURRENT(self, gcmd):
-        prev_run_current, prev_hold_current, max_current = self.get_current()
-        run_current = gcmd.get_float('CURRENT', None,
-                                     minval=0., maxval=max_current)
-        hold_current = gcmd.get_float('HOLDCURRENT', None,
-                                      above=0., maxval=max_current)
-        if run_current is None and hold_current is None:
-            # Query only
-            gcmd.respond_info("Run Current: %0.2fA Hold Current: %0.2fA"
-                              % (prev_run_current, prev_hold_current))
-            return
-        if run_current is None:
-            run_current = prev_run_current
-        if hold_current is None:
-            hold_current = prev_hold_current
-        print_time = self.printer.lookup_object('toolhead').get_last_move_time()
-        self.set_current(run_current, hold_current, print_time)
 
 
 ######################################################################
@@ -261,10 +239,10 @@ class TMC2130:
         # Allow virtual pins to be created
         tmc.TMCVirtualPinHelper(config, self.mcu_tmc)
         # Register commands
-        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
+        current_helper = TMCCurrentHelper(config, self.mcu_tmc)
+        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc, current_helper)
         cmdhelper.setup_register_dump(ReadRegisters)
         # Setup basic register values
-        TMCCurrentHelper(config, self.mcu_tmc)
         mh = tmc.TMCMicrostepHelper(config, self.mcu_tmc)
         self.get_microsteps = mh.get_microsteps
         self.get_phase = mh.get_phase
diff --git a/klippy/extras/tmc2208.py b/klippy/extras/tmc2208.py
index 31368e3c..5127d61e 100644
--- a/klippy/extras/tmc2208.py
+++ b/klippy/extras/tmc2208.py
@@ -188,13 +188,13 @@ class TMC2208:
         self.fields = tmc.FieldHelper(Fields, SignedFields, FieldFormatters)
         self.mcu_tmc = tmc_uart.MCU_TMC_uart(config, Registers, self.fields)
         # Register commands
-        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
+        current_helper = tmc2130.TMCCurrentHelper(config, self.mcu_tmc)
+        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc, current_helper)
         cmdhelper.setup_register_dump(ReadRegisters, self.read_translate)
         # Setup basic register values
         self.fields.set_field("pdn_disable", True)
         self.fields.set_field("mstep_reg_select", True)
         self.fields.set_field("multistep_filt", True)
-        tmc2130.TMCCurrentHelper(config, self.mcu_tmc)
         mh = tmc.TMCMicrostepHelper(config, self.mcu_tmc)
         self.get_microsteps = mh.get_microsteps
         self.get_phase = mh.get_phase
diff --git a/klippy/extras/tmc2209.py b/klippy/extras/tmc2209.py
index 666bba63..ceb616e9 100644
--- a/klippy/extras/tmc2209.py
+++ b/klippy/extras/tmc2209.py
@@ -62,13 +62,13 @@ class TMC2209:
         # Allow virtual pins to be created
         tmc.TMCVirtualPinHelper(config, self.mcu_tmc)
         # Register commands
-        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
+        current_helper = tmc2130.TMCCurrentHelper(config, self.mcu_tmc)
+        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc, current_helper)
         cmdhelper.setup_register_dump(ReadRegisters)
         # Setup basic register values
         self.fields.set_field("pdn_disable", True)
         self.fields.set_field("mstep_reg_select", True)
         self.fields.set_field("multistep_filt", True)
-        tmc2130.TMCCurrentHelper(config, self.mcu_tmc)
         mh = tmc.TMCMicrostepHelper(config, self.mcu_tmc)
         self.get_microsteps = mh.get_microsteps
         self.get_phase = mh.get_phase
diff --git a/klippy/extras/tmc2660.py b/klippy/extras/tmc2660.py
index 129974c4..44280c86 100644
--- a/klippy/extras/tmc2660.py
+++ b/klippy/extras/tmc2660.py
@@ -149,11 +149,6 @@ class TMC2660CurrentHelper:
             self.printer.register_event_handler("idle_timeout:ready",
                                                 self._handle_ready)
 
-        gcode = self.printer.lookup_object("gcode")
-        gcode.register_mux_command("SET_TMC_CURRENT", "STEPPER", self.name,
-                                   self.cmd_SET_TMC_CURRENT,
-                                   desc=self.cmd_SET_TMC_CURRENT_help)
-
     def _calc_current_bits(self, current, vsense):
         vref = 0.165 if vsense else 0.310
         cs = int(32 * current * self.sense_resistor * math.sqrt(2.) / vref
@@ -194,14 +189,6 @@ class TMC2660CurrentHelper:
         self.current = run_current
         self._update_current(run_current, print_time)
 
-    cmd_SET_TMC_CURRENT_help = "Set the current of a TMC2660 driver"
-    def cmd_SET_TMC_CURRENT(self, gcmd):
-        cur = gcmd.get_float('CURRENT', None, minval=0.1, maxval=MAX_CURRENT)
-        if cur is None:
-            return
-        print_time = self.printer.lookup_object('toolhead').get_last_move_time()
-        self.set_current(cur, None, print_time)
-
 
 ######################################################################
 # TMC2660 SPI
@@ -248,7 +235,8 @@ class TMC2660:
         self.fields.set_field("SDOFF", 0) # Access DRVCTRL in step/dir mode
         self.mcu_tmc = MCU_TMC2660_SPI(config, Registers, self.fields)
         # Register commands
-        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
+        current_helper = TMC2660CurrentHelper(config, self.mcu_tmc)
+        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc, current_helper)
         cmdhelper.setup_register_dump(ReadRegisters)
 
         # DRVCTRL
@@ -278,7 +266,6 @@ class TMC2660:
         # SGSCONF
         set_config_field(config, "SFILT", 1)
         set_config_field(config, "SGT", 0)
-        TMC2660CurrentHelper(config, self.mcu_tmc)
 
         # DRVCONF
         set_config_field(config, "SLPH", 0)
diff --git a/klippy/extras/tmc5160.py b/klippy/extras/tmc5160.py
index 4d6f921a..4447bd75 100644
--- a/klippy/extras/tmc5160.py
+++ b/klippy/extras/tmc5160.py
@@ -244,10 +244,6 @@ class TMC5160CurrentHelper:
         irun, ihold = self._calc_current(run_current, hold_current)
         self.fields.set_field("IHOLD", ihold)
         self.fields.set_field("IRUN", irun)
-        gcode = self.printer.lookup_object("gcode")
-        gcode.register_mux_command("SET_TMC_CURRENT", "STEPPER", self.name,
-                                   self.cmd_SET_TMC_CURRENT,
-                                   desc=self.cmd_SET_TMC_CURRENT_help)
     def _set_globalscaler(self, current):
         globalscaler = int((current * 256. * math.sqrt(2.)
                             * self.sense_resistor / VREF) + .5)
@@ -283,24 +279,6 @@ class TMC5160CurrentHelper:
         self.fields.set_field("IHOLD", ihold)
         val = self.fields.set_field("IRUN", irun)
         self.mcu_tmc.set_register("IHOLD_IRUN", val, print_time)
-    cmd_SET_TMC_CURRENT_help = "Set the current of a TMC driver"
-    def cmd_SET_TMC_CURRENT(self, gcmd):
-        prev_run_current, prev_hold_current, max_current = self.get_current()
-        run_current = gcmd.get_float('CURRENT', None,
-                                     minval=0., maxval=max_current)
-        hold_current = gcmd.get_float('HOLDCURRENT', None,
-                                      above=0., maxval=max_current)
-        if run_current is None and hold_current is None:
-            # Query only
-            gcmd.respond_info("Run Current: %0.2fA Hold Current: %0.2fA"
-                              % (prev_run_current, prev_hold_current))
-            return
-        if run_current is None:
-            run_current = prev_run_current
-        if hold_current is None:
-            hold_current = prev_hold_current
-        print_time = self.printer.lookup_object('toolhead').get_last_move_time()
-        self.set_current(run_current, hold_current, print_time)
 
 
 ######################################################################
@@ -315,7 +293,8 @@ class TMC5160:
         # Allow virtual pins to be created
         tmc.TMCVirtualPinHelper(config, self.mcu_tmc)
         # Register commands
-        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
+        current_helper = TMC5160CurrentHelper(config, self.mcu_tmc)
+        cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc, current_helper)
         cmdhelper.setup_register_dump(ReadRegisters)
         # Setup basic register values
         mh = tmc.TMCMicrostepHelper(config, self.mcu_tmc)
@@ -345,7 +324,6 @@ class TMC5160:
         set_config_field(config, "sgt", 0)
         set_config_field(config, "sfilt", 0)
         #   IHOLDIRUN
-        TMC5160CurrentHelper(config, self.mcu_tmc)
         set_config_field(config, "IHOLDDELAY", 6)
         #   PWMCONF
         set_config_field(config, "PWM_OFS", 30)