resonance_tester: Support CHIPS= parameter in SHAPER_CALIBRATE command

Signed-off-by: Dmitry Butyugin <dmbutyugin@google.com>

2 files changed, 18 insertions, 14 deletions

diff --git a/docs/G-Codes.md b/docs/G-Codes.md
index d89c7df8..fe24a580 100644
--- a/docs/G-Codes.md
+++ b/docs/G-Codes.md
@@ -1018,7 +1018,7 @@ frequency response is calculated (across all probe points) and written into
 
 #### SHAPER_CALIBRATE
 `SHAPER_CALIBRATE [AXIS=<axis>] [NAME=<name>] [FREQ_START=<min_freq>]
-[FREQ_END=<max_freq>] [HZ_PER_SEC=<hz_per_sec>]
+[FREQ_END=<max_freq>] [HZ_PER_SEC=<hz_per_sec>] [CHIPS=<adxl345_chip_name>]
 [MAX_SMOOTHING=<max_smoothing>]`: Similarly to `TEST_RESONANCES`, runs
 the resonance test as configured, and tries to find the optimal
 parameters for the input shaper for the requested axis (or both X and
diff --git a/klippy/extras/resonance_tester.py b/klippy/extras/resonance_tester.py
index db03e320..819d8155 100644
--- a/klippy/extras/resonance_tester.py
+++ b/klippy/extras/resonance_tester.py
@@ -207,11 +207,21 @@ class ResonanceTester:
                     else:
                         calibration_data[axis].add_data(new_data)
         return calibration_data
+    def _parse_chips(self, accel_chips):
+        parsed_chips = []
+        for chip_name in accel_chips.split(','):
+            if "adxl345" in chip_name:
+                chip_lookup_name = chip_name.strip()
+            else:
+                chip_lookup_name = "adxl345 " + chip_name.strip();
+            chip = self.printer.lookup_object(chip_lookup_name)
+            parsed_chips.append(chip)
+        return parsed_chips
     cmd_TEST_RESONANCES_help = ("Runs the resonance test for a specifed axis")
     def cmd_TEST_RESONANCES(self, gcmd):
         # Parse parameters
         axis = _parse_axis(gcmd, gcmd.get("AXIS").lower())
-        accel_chips = gcmd.get("CHIPS", None)
+        chips_str = gcmd.get("CHIPS", None)
         test_point = gcmd.get("POINT", None)
 
         if test_point:
@@ -224,15 +234,7 @@ class ResonanceTester:
                 raise gcmd.error("Invalid POINT parameter, must be 'x,y,z'"
                 " where x, y and z are valid floating point numbers")
 
-        if accel_chips:
-            parsed_chips = []
-            for chip_name in accel_chips.split(','):
-                if "adxl345" in chip_name:
-                    chip_lookup_name = chip_name.strip()
-                else:
-                    chip_lookup_name = "adxl345 " + chip_name.strip();
-                chip = self.printer.lookup_object(chip_lookup_name)
-                parsed_chips.append(chip)
+        accel_chips = self._parse_chips(chips_str) if chips_str else None
 
         outputs = gcmd.get("OUTPUT", "resonances").lower().split(',')
         for output in outputs:
@@ -257,8 +259,7 @@ class ResonanceTester:
         data = self._run_test(
                 gcmd, [axis], helper,
                 raw_name_suffix=name_suffix if raw_output else None,
-                accel_chips=parsed_chips if accel_chips else None,
-                test_point=test_point)[axis]
+                accel_chips=accel_chips, test_point=test_point)[axis]
         if csv_output:
             csv_name = self.save_calibration_data('resonances', name_suffix,
                                                   helper, axis, data,
@@ -276,6 +277,8 @@ class ResonanceTester:
             raise gcmd.error("Unsupported axis '%s'" % (axis,))
         else:
             calibrate_axes = [TestAxis(axis.lower())]
+        chips_str = gcmd.get("CHIPS", None)
+        accel_chips = self._parse_chips(chips_str) if chips_str else None
 
         max_smoothing = gcmd.get_float(
                 "MAX_SMOOTHING", self.max_smoothing, minval=0.05)
@@ -287,7 +290,8 @@ class ResonanceTester:
         # Setup shaper calibration
         helper = shaper_calibrate.ShaperCalibrate(self.printer)
 
-        calibration_data = self._run_test(gcmd, calibrate_axes, helper)
+        calibration_data = self._run_test(gcmd, calibrate_axes, helper,
+                                          accel_chips=accel_chips)
 
         configfile = self.printer.lookup_object('configfile')
         for axis in calibrate_axes: