tmc: Simplify internal DUMP_TMC interface

Only the tmc2208 driver needs special read register translation logic.
Rework the code so a default implementation is available for all the
other drivers.

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

6 files changed, 19 insertions, 36 deletions

diff --git a/klippy/extras/tmc.py b/klippy/extras/tmc.py
index d0c88f58..46b9a8cf 100644
--- a/klippy/extras/tmc.py
+++ b/klippy/extras/tmc.py
@@ -86,7 +86,7 @@ class TMCCommandHelper:
         self.name = config.get_name().split()[-1]
         self.mcu_tmc = mcu_tmc
         self.fields = mcu_tmc.get_fields()
-        self.query_registers = None
+        self.read_registers = self.read_translate = None
         self.gcode = self.printer.lookup_object("gcode")
         self.gcode.register_mux_command(
             "SET_TMC_FIELD", "STEPPER", self.name,
@@ -131,8 +131,9 @@ class TMCCommandHelper:
         print_time = self.printer.lookup_object('toolhead').get_last_move_time()
         self.mcu_tmc.set_register(reg_name, reg_val, print_time)
     # DUMP_TMC support
-    def setup_register_dump(self, query_registers):
-        self.query_registers = query_registers
+    def setup_register_dump(self, read_registers, read_translate=None):
+        self.read_registers = read_registers
+        self.read_translate = read_translate
         self.gcode.register_mux_command(
             "DUMP_TMC", "STEPPER", self.name,
             self.cmd_DUMP_TMC, desc=self.cmd_DUMP_TMC_help)
@@ -140,15 +141,16 @@ class TMCCommandHelper:
     def cmd_DUMP_TMC(self, params):
         logging.info("DUMP_TMC %s", self.name)
         print_time = self.printer.lookup_object('toolhead').get_last_move_time()
-        read_regs = self.query_registers(print_time)
-        read_regs_by_name = { reg_name: val for reg_name, val in read_regs }
         self.gcode.respond_info("========== Write-only registers ==========")
         for reg_name, val in self.fields.registers.items():
-            if reg_name not in read_regs_by_name:
+            if reg_name not in self.read_registers:
                 self.gcode.respond_info(
                     self.fields.pretty_format(reg_name, val))
         self.gcode.respond_info("========== Queried registers ==========")
-        for reg_name, val in read_regs:
+        for reg_name in self.read_registers:
+            val = self.mcu_tmc.get_register(reg_name)
+            if self.read_translate is not None:
+                reg_name, val = self.read_translate(reg_name, val)
             self.gcode.respond_info(self.fields.pretty_format(reg_name, val))
 
 
diff --git a/klippy/extras/tmc2130.py b/klippy/extras/tmc2130.py
index aae4b2b1..4dce19c9 100644
--- a/klippy/extras/tmc2130.py
+++ b/klippy/extras/tmc2130.py
@@ -208,7 +208,7 @@ class TMC2130:
         tmc.TMCEndstopHelper(config, self.mcu_tmc, diag1_pin)
         # Register commands
         cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
-        cmdhelper.setup_register_dump(self.query_registers)
+        cmdhelper.setup_register_dump(ReadRegisters)
         # Setup basic register values
         TMCCurrentHelper(config, self.mcu_tmc)
         mh = tmc.TMCMicrostepHelper(config, self.mcu_tmc)
@@ -229,9 +229,6 @@ class TMC2130:
         set_config_field(config, "pwm_freq", 1)
         set_config_field(config, "pwm_autoscale", True)
         set_config_field(config, "sgt", 0)
-    def query_registers(self, print_time=0.):
-        return [(reg_name, self.mcu_tmc.get_register(reg_name))
-                for reg_name in ReadRegisters]
 
 def load_config_prefix(config):
     return TMC2130(config)
diff --git a/klippy/extras/tmc2208.py b/klippy/extras/tmc2208.py
index 6ea0c056..4bcdabd8 100644
--- a/klippy/extras/tmc2208.py
+++ b/klippy/extras/tmc2208.py
@@ -189,7 +189,7 @@ class TMC2208:
         self.mcu_tmc = tmc_uart.MCU_TMC_uart(config, Registers, self.fields)
         # Register commands
         cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
-        cmdhelper.setup_register_dump(self.query_registers)
+        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)
@@ -215,17 +215,11 @@ class TMC2208:
         set_config_field(config, "pwm_autograd", True)
         set_config_field(config, "PWM_REG", 8)
         set_config_field(config, "PWM_LIM", 12)
-    def query_registers(self, print_time=0.):
-        out = []
-        for reg_name in ReadRegisters:
-            val = self.mcu_tmc.get_register(reg_name)
-            # IOIN has different mappings depending on the driver type
-            # (SEL_A field of IOIN reg)
-            if reg_name == "IOIN":
-                drv_type = self.fields.get_field("SEL_A", val)
-                reg_name = "IOIN@TMC220x" if drv_type else "IOIN@TMC222x"
-            out.append((reg_name, val))
-        return out
+    def read_translate(self, reg_name, val):
+        if reg_name == "IOIN":
+            drv_type = self.fields.get_field("SEL_A", val)
+            reg_name = "IOIN@TMC220x" if drv_type else "IOIN@TMC222x"
+        return reg_name, val
 
 def load_config_prefix(config):
     return TMC2208(config)
diff --git a/klippy/extras/tmc2209.py b/klippy/extras/tmc2209.py
index 1e012153..22a973ff 100644
--- a/klippy/extras/tmc2209.py
+++ b/klippy/extras/tmc2209.py
@@ -61,7 +61,7 @@ class TMC2209:
         self.mcu_tmc = tmc_uart.MCU_TMC_uart(config, Registers, self.fields)
         # Register commands
         cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
-        cmdhelper.setup_register_dump(self.query_registers)
+        cmdhelper.setup_register_dump(ReadRegisters)
         # Setup basic register values
         self.fields.set_field("pdn_disable", True)
         self.fields.set_field("mstep_reg_select", True)
@@ -88,9 +88,6 @@ class TMC2209:
         set_config_field(config, "PWM_REG", 8)
         set_config_field(config, "PWM_LIM", 12)
         set_config_field(config, "SGTHRS", 0)
-    def query_registers(self, print_time=0.):
-        return [(reg_name, self.mcu_tmc.get_register(reg_name))
-                for reg_name in ReadRegisters]
 
 def load_config_prefix(config):
     return TMC2209(config)
diff --git a/klippy/extras/tmc2660.py b/klippy/extras/tmc2660.py
index c894ed36..893f8d1b 100644
--- a/klippy/extras/tmc2660.py
+++ b/klippy/extras/tmc2660.py
@@ -236,7 +236,7 @@ class TMC2660:
         self.mcu_tmc = MCU_TMC2660_SPI(config, Registers, self.fields)
         # Register commands
         cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
-        cmdhelper.setup_register_dump(self.query_registers)
+        cmdhelper.setup_register_dump(ReadRegisters)
 
         # DRVCTRL
         mh = tmc.TMCMicrostepHelper(config, self.mcu_tmc)
@@ -274,9 +274,5 @@ class TMC2660:
         set_config_field(config, "DISS2G", 0)
         set_config_field(config, "TS2G", 3)
 
-    def query_registers(self, print_time=0.):
-        return [(reg_name, self.mcu_tmc.get_register(reg_name))
-                for reg_name in ReadRegisters]
-
 def load_config_prefix(config):
     return TMC2660(config)
diff --git a/klippy/extras/tmc5160.py b/klippy/extras/tmc5160.py
index d9926638..637f8d39 100644
--- a/klippy/extras/tmc5160.py
+++ b/klippy/extras/tmc5160.py
@@ -296,7 +296,7 @@ class TMC5160:
         tmc.TMCEndstopHelper(config, self.mcu_tmc, diag1_pin)
         # Register commands
         cmdhelper = tmc.TMCCommandHelper(config, self.mcu_tmc)
-        cmdhelper.setup_register_dump(self.query_registers)
+        cmdhelper.setup_register_dump(ReadRegisters)
         # Setup basic register values
         mh = tmc.TMCMicrostepHelper(config, self.mcu_tmc)
         self.get_microsteps = mh.get_microsteps
@@ -339,9 +339,6 @@ class TMC5160:
         set_config_field(config, "PWM_LIM", 12)
         #   TPOWERDOWN
         set_config_field(config, "TPOWERDOWN", 10)
-    def query_registers(self, print_time=0.):
-        return [(reg_name, self.mcu_tmc.get_register(reg_name))
-                for reg_name in ReadRegisters]
 
 def load_config_prefix(config):
     return TMC5160(config)