Reduce code duplication (#182)

Frix-x · Dec 29, 2024 · f07809f · f07809f
1 parent 473631d
commit f07809f
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 32 deletions.
diff --git a/shaketune/commands/axes_shaper_calibration.py b/shaketune/commands/axes_shaper_calibration.py
@@ -11,6 +11,7 @@
 
 from ..helpers.accelerometer import Accelerometer, MeasurementsManager
 from ..helpers.common_func import AXIS_CONFIG
+from ..helpers.compat import res_tester_config
 from ..helpers.console_output import ConsoleOutput
 from ..helpers.resonance_test import vibrate_axis
 from ..shaketune_process import ShakeTuneProcess
@@ -24,18 +25,7 @@ def axes_shaper_calibration(gcmd, config, st_process: ShakeTuneProcess) -> None:
     toolhead_info = toolhead.get_status(systime)
 
     # Get the default values for the frequency range and the acceleration per Hz
-    if hasattr(res_tester, 'test'):
-        # Old Klipper code (before Dec 6, 2024: https://github.com/Klipper3d/klipper/commit/16b4b6b302ac3ffcd55006cd76265aad4e26ecc8)
-        default_min_freq = res_tester.test.min_freq
-        default_max_freq = res_tester.test.max_freq
-        default_accel_per_hz = res_tester.test.accel_per_hz
-        test_points = res_tester.test.get_start_test_points()
-    else:
-        # New Klipper code (after Dec 6, 2024) with the sweeping test
-        default_min_freq = res_tester.generator.vibration_generator.min_freq
-        default_max_freq = res_tester.generator.vibration_generator.max_freq
-        default_accel_per_hz = res_tester.generator.vibration_generator.accel_per_hz
-        test_points = res_tester.probe_points
+    default_min_freq, default_max_freq, default_accel_per_hz, test_points = res_tester_config(config)
 
     min_freq = gcmd.get_float('FREQ_START', default=default_min_freq, minval=1)
     max_freq = gcmd.get_float('FREQ_END', default=default_max_freq, minval=1)

diff --git a/shaketune/commands/compare_belts_responses.py b/shaketune/commands/compare_belts_responses.py
@@ -11,6 +11,7 @@
 
 from ..helpers.accelerometer import Accelerometer, MeasurementsManager
 from ..helpers.common_func import AXIS_CONFIG
+from ..helpers.compat import res_tester_config
 from ..helpers.console_output import ConsoleOutput
 from ..helpers.motors_config_parser import MotorsConfigParser
 from ..helpers.resonance_test import vibrate_axis
@@ -24,18 +25,7 @@ def compare_belts_responses(gcmd, config, st_process: ShakeTuneProcess) -> None:
     systime = printer.get_reactor().monotonic()
 
     # Get the default values for the frequency range and the acceleration per Hz
-    if hasattr(res_tester, 'test'):
-        # Old Klipper code (before Dec 6, 2024: https://github.com/Klipper3d/klipper/commit/16b4b6b302ac3ffcd55006cd76265aad4e26ecc8)
-        default_min_freq = res_tester.test.min_freq
-        default_max_freq = res_tester.test.max_freq
-        default_accel_per_hz = res_tester.test.accel_per_hz
-        test_points = res_tester.test.get_start_test_points()
-    else:
-        # New Klipper code (after Dec 6, 2024) with the sweeping test
-        default_min_freq = res_tester.generator.vibration_generator.min_freq
-        default_max_freq = res_tester.generator.vibration_generator.max_freq
-        default_accel_per_hz = res_tester.generator.vibration_generator.accel_per_hz
-        test_points = res_tester.probe_points
+    default_min_freq, default_max_freq, default_accel_per_hz, test_points = res_tester_config(config)
 
     min_freq = gcmd.get_float('FREQ_START', default=default_min_freq, minval=1)
     max_freq = gcmd.get_float('FREQ_END', default=default_max_freq, minval=1)

diff --git a/shaketune/commands/excitate_axis_at_freq.py b/shaketune/commands/excitate_axis_at_freq.py
@@ -10,6 +10,7 @@
 
 from ..helpers.accelerometer import Accelerometer, MeasurementsManager
 from ..helpers.common_func import AXIS_CONFIG
+from ..helpers.compat import res_tester_config
 from ..helpers.console_output import ConsoleOutput
 from ..helpers.resonance_test import vibrate_axis_at_static_freq
 from ..shaketune_process import ShakeTuneProcess
@@ -53,14 +54,7 @@ def excitate_axis_at_freq(gcmd, config, st_process: ShakeTuneProcess) -> None:
     systime = printer.get_reactor().monotonic()
 
     # Get the default values for the acceleration per Hz and the test points
-    if hasattr(res_tester, 'test'):
-        # Old Klipper code (before Dec 6, 2024: https://github.com/Klipper3d/klipper/commit/16b4b6b302ac3ffcd55006cd76265aad4e26ecc8)
-        default_accel_per_hz = res_tester.test.accel_per_hz
-        test_points = res_tester.test.get_start_test_points()
-    else:
-        # New Klipper code (after Dec 6, 2024) with the sweeping test
-        default_accel_per_hz = res_tester.generator.vibration_generator.accel_per_hz
-        test_points = res_tester.probe_points
+    default_min_freq, default_max_freq, default_accel_per_hz, test_points = res_tester_config(config)
 
     if accel_per_hz is None:
         accel_per_hz = default_accel_per_hz

diff --git a/shaketune/helpers/compat.py b/shaketune/helpers/compat.py
@@ -0,0 +1,31 @@
+# Shake&Tune: 3D printer analysis tools
+#
+# Copyright (C) 2024 Félix Boisselier <[email protected]> (Frix_x on Discord)
+# Licensed under the GNU General Public License v3.0 (GPL-3.0)
+#
+# File: compat.py
+# Description: Handles compatibility with different versions of Klipper.
+
+from collections import namedtuple
+
+ResTesterConfig = namedtuple('ResTesterConfig', ['default_min_freq', 'default_max_freq', 'default_accel_per_hz', 'test_points'])
+
+def res_tester_config(config) -> ResTesterConfig:
+    printer = config.get_printer()
+    res_tester = printer.lookup_object('resonance_tester')
+
+    # Get the default values for the frequency range and the acceleration per Hz
+    if hasattr(res_tester, 'test'):
+        # Old Klipper code (before Dec 6, 2024: https://github.com/Klipper3d/klipper/commit/16b4b6b302ac3ffcd55006cd76265aad4e26ecc8)
+        default_min_freq = res_tester.test.min_freq
+        default_max_freq = res_tester.test.max_freq
+        default_accel_per_hz = res_tester.test.accel_per_hz
+        test_points = res_tester.test.get_start_test_points()
+    else:
+        # New Klipper code (after Dec 6, 2024) with the sweeping test
+        default_min_freq = res_tester.generator.vibration_generator.min_freq
+        default_max_freq = res_tester.generator.vibration_generator.max_freq
+        default_accel_per_hz = res_tester.generator.vibration_generator.accel_per_hz
+        test_points = res_tester.probe_points
+
+    return ResTesterConfig(default_min_freq, default_max_freq, default_accel_per_hz, test_points)