tmc2130: Add set_current()/get_current() helper functions to TMCCurrentHelper

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

1 files changed, 20 insertions, 16 deletions

diff --git a/klippy/extras/tmc2130.py b/klippy/extras/tmc2130.py
index f300afbc..0f503726 100644
--- a/klippy/extras/tmc2130.py
+++ b/klippy/extras/tmc2130.py
@@ -137,33 +137,37 @@ class TMCCurrentHelper:
         vref = 0.32
         if self.fields.get_field("vsense"):
             vref = 0.18
-        current = (bits + 1) * vref / (32 * sense_resistor * math.sqrt(2.))
-        return round(current, 2)
+        return (bits + 1) * vref / (32 * sense_resistor * math.sqrt(2.))
+    def get_current(self):
+        run_current = self._calc_current_from_field("IRUN")
+        hold_current = self._calc_current_from_field("IHOLD")
+        return run_current, hold_current, MAX_CURRENT
+    def set_current(self, run_current, hold_current, print_time):
+        vsense, irun, ihold = self._calc_current(run_current, hold_current)
+        if vsense != self.fields.get_field("vsense"):
+            val = self.fields.set_field("vsense", vsense)
+            self.mcu_tmc.set_register("CHOPCONF", val, print_time)
+        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)
+                                     minval=0., maxval=max_current)
         hold_current = gcmd.get_float('HOLDCURRENT', None,
-                                      above=0., maxval=MAX_CURRENT)
+                                      above=0., maxval=max_current)
         if run_current is None and hold_current is None:
             # Query only
-            run_current = self._calc_current_from_field("IRUN")
-            hold_current = self._calc_current_from_field("IHOLD")
             gcmd.respond_info("Run Current: %0.2fA Hold Current: %0.2fA"
-                              % (run_current, hold_current))
+                              % (prev_run_current, prev_hold_current))
             return
         if run_current is None:
-            run_current = self._calc_current_from_field("IRUN")
+            run_current = prev_run_current
         if hold_current is None:
-            hold_current = self._calc_current_from_field("IHOLD")
+            hold_current = prev_hold_current
         print_time = self.printer.lookup_object('toolhead').get_last_move_time()
-        vsense, irun, ihold = self._calc_current(run_current, hold_current)
-        if vsense != self.fields.get_field("vsense"):
-            val = self.fields.set_field("vsense", vsense)
-            self.mcu_tmc.set_register("CHOPCONF", val, print_time)
-        self.fields.set_field("IHOLD", ihold)
-        val = self.fields.set_field("IRUN", irun)
-        self.mcu_tmc.set_register("IHOLD_IRUN", val, print_time)
+        self.set_current(run_current, hold_current, print_time)
 
 
 ######################################################################