implemented a Factory design pattern for graph creators

shaketune/graph_creators/__init__.py

@@ -10,6 +10,7 @@
 from .axes_map_graph_creator import AxesMapGraphCreator as AxesMapGraphCreator
 from .belts_graph_creator import BeltsGraphCreator as BeltsGraphCreator
 from .graph_creator import GraphCreator as GraphCreator
+from .graph_creator_factory import GraphCreatorFactory as GraphCreatorFactory
 from .shaper_graph_creator import ShaperGraphCreator as ShaperGraphCreator
 from .static_graph_creator import StaticGraphCreator as StaticGraphCreator
 from .vibrations_graph_creator import VibrationsGraphCreator as VibrationsGraphCreator

shaketune/graph_creators/axes_map_graph_creator.py

@@ -38,9 +38,10 @@
 MACHINE_AXES = ['x', 'y', 'z']
 
 
+@GraphCreator.register('axes map')
 class AxesMapGraphCreator(GraphCreator):
     def __init__(self, config: ShakeTuneConfig):
-        super().__init__(config, 'axes map')
+        super().__init__(config)
         self._accel: Optional[int] = None
         self._segment_length: Optional[float] = None
 

shaketune/graph_creators/belts_graph_creator.py

@@ -62,9 +62,10 @@ class PeakPairingResult(NamedTuple):
     unpaired_peaks2: List[int]
 
 
+@GraphCreator.register('belts comparison')
 class BeltsGraphCreator(GraphCreator):
     def __init__(self, config: ShakeTuneConfig):
-        super().__init__(config, 'belts comparison')
+        super().__init__(config)
         self._kinematics: Optional[str] = None
         self._accel_per_hz: Optional[float] = None
 

shaketune/graph_creators/graph_creator.py

@@ -21,12 +21,23 @@
 
 
 class GraphCreator(abc.ABC):
-    def __init__(self, config: ShakeTuneConfig, graph_type: str):
+    registry = {}
+
+    @classmethod
+    def register(cls, graph_type: str):
+        def decorator(subclass):
+            cls.registry[graph_type] = subclass
+            subclass.graph_type = graph_type
+            return subclass
+
+        return decorator
+
+    def __init__(self, config: ShakeTuneConfig):
         self._config = config
         self._graph_date = datetime.now().strftime('%Y%m%d_%H%M%S')
         self._version = ShakeTuneConfig.get_git_version()
-        self._type = graph_type
-        self._folder = self._config.get_results_folder(graph_type)
+        self._type = self.__class__.graph_type
+        self._folder = self._config.get_results_folder(self._type)
 
     def _save_figure(
         self, fig: Figure, measurements_manager: MeasurementsManager, axis_label: Optional[str] = None

shaketune/graph_creators/graph_creator_factory.py

@@ -0,0 +1,20 @@
+# 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: graph_creator_factory.py
+# Description: Contains a factory class to create the different graph creators.
+
+
+from ..shaketune_config import ShakeTuneConfig
+from .graph_creator import GraphCreator
+
+
+class GraphCreatorFactory:
+    @staticmethod
+    def create_graph_creator(graph_type: str, config: ShakeTuneConfig) -> GraphCreator:
+        creator_class = GraphCreator.registry.get(graph_type, None)
+        if not creator_class:
+            raise NotImplementedError(f'Graph creator for {graph_type} not implemented!')
+        return creator_class(config)

shaketune/graph_creators/shaper_graph_creator.py

@@ -60,9 +60,10 @@
 }
 
 
+@GraphCreator.register('input shaper')
 class ShaperGraphCreator(GraphCreator):
     def __init__(self, config: ShakeTuneConfig):
-        super().__init__(config, 'input shaper')
+        super().__init__(config)
         self._max_smoothing: Optional[float] = None
         self._scv: Optional[float] = None
         self._accel_per_hz: Optional[float] = None

shaketune/graph_creators/static_graph_creator.py

@@ -41,9 +41,10 @@
 }
 
 
+@GraphCreator.register('static frequency')
 class StaticGraphCreator(GraphCreator):
     def __init__(self, config: ShakeTuneConfig):
-        super().__init__(config, 'static frequency')
+        super().__init__(config)
         self._freq: Optional[float] = None
         self._duration: Optional[float] = None
         self._accel_per_hz: Optional[float] = None

shaketune/graph_creators/vibrations_graph_creator.py

@@ -52,9 +52,10 @@
 }
 
 
+@GraphCreator.register('vibrations profile')
 class VibrationsGraphCreator(GraphCreator):
     def __init__(self, config: ShakeTuneConfig):
-        super().__init__(config, 'vibrations profile')
+        super().__init__(config)
         self._kinematics: Optional[str] = None
         self._accel: Optional[float] = None
         self._motors: Optional[List[MotorsConfigParser]] = None
@@ -848,7 +849,7 @@ def vibrations_profile(
     if motors is not None and len(motors) == 2:
         differences = motors[0].compare_to(motors[1])
         plot_motor_config_txt(fig, motors, differences)
-        if differences is not None and kinematics in {'corexy', 'limited_corexy'}: 
+        if differences is not None and kinematics in {'corexy', 'limited_corexy'}:
             ConsoleOutput.print(f'Warning: motors have different TMC configurations!\n{differences}')
 
     # Plot the graphs
@@ -908,8 +909,17 @@ def main():
         opts.error('No measurements to analyse')
     if options.output is None:
         opts.error('You must specify an output file.png to use the script (option -o)')
-    if options.kinematics not in {'cartesian', 'limited_cartesian', 'corexy', 'limited_corexy', 'corexz', 'limited_corexz'}:
-        opts.error('Only cartesian, limited_cartesian, corexy, limited_corexy, corexz, limited_corexz kinematics are supported by this tool at the moment!')
+    if options.kinematics not in {
+        'cartesian',
+        'limited_cartesian',
+        'corexy',
+        'limited_corexy',
+        'corexz',
+        'limited_corexz',
+    }:
+        opts.error(
+            'Only cartesian, limited_cartesian, corexy, limited_corexy, corexz, limited_corexz kinematics are supported by this tool at the moment!'
+        )
 
     measurements_manager = MeasurementsManager(10)
     if args[0].endswith('.csv'):

shaketune/shaketune.py

@@ -19,13 +19,7 @@
     create_vibrations_profile,
     excitate_axis_at_freq,
 )
-from .graph_creators import (
-    AxesMapGraphCreator,
-    BeltsGraphCreator,
-    ShaperGraphCreator,
-    StaticGraphCreator,
-    VibrationsGraphCreator,
-)
+from .graph_creators import GraphCreatorFactory
 from .helpers.console_output import ConsoleOutput
 from .shaketune_config import ShakeTuneConfig
 from .shaketune_process import ShakeTuneProcess
@@ -153,55 +147,55 @@ def _on_klippy_connect(self) -> None:
 
     def cmd_EXCITATE_AXIS_AT_FREQ(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
-        static_freq_graph_creator = StaticGraphCreator(self._st_config)
+        gcreator = GraphCreatorFactory.create_graph_creator('static frequency', self._st_config)
         st_process = ShakeTuneProcess(
             self._st_config,
             self._printer.get_reactor(),
-            static_freq_graph_creator,
+            gcreator,
             self.timeout,
         )
         excitate_axis_at_freq(gcmd, self._config, st_process)
 
     def cmd_AXES_MAP_CALIBRATION(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
-        axes_map_graph_creator = AxesMapGraphCreator(self._st_config)
+        gcreator = GraphCreatorFactory.create_graph_creator('axes map', self._st_config)
         st_process = ShakeTuneProcess(
             self._st_config,
             self._printer.get_reactor(),
-            axes_map_graph_creator,
+            gcreator,
             self.timeout,
         )
         axes_map_calibration(gcmd, self._config, st_process)
 
     def cmd_COMPARE_BELTS_RESPONSES(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
-        belt_graph_creator = BeltsGraphCreator(self._st_config)
+        gcreator = GraphCreatorFactory.create_graph_creator('belts comparison', self._st_config)
         st_process = ShakeTuneProcess(
             self._st_config,
             self._printer.get_reactor(),
-            belt_graph_creator,
+            gcreator,
             self.timeout,
         )
         compare_belts_responses(gcmd, self._config, st_process)
 
     def cmd_AXES_SHAPER_CALIBRATION(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
-        shaper_graph_creator = ShaperGraphCreator(self._st_config)
+        gcreator = GraphCreatorFactory.create_graph_creator('input shaper', self._st_config)
         st_process = ShakeTuneProcess(
             self._st_config,
             self._printer.get_reactor(),
-            shaper_graph_creator,
+            gcreator,
             self.timeout,
         )
         axes_shaper_calibration(gcmd, self._config, st_process)
 
     def cmd_CREATE_VIBRATIONS_PROFILE(self, gcmd) -> None:
         ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
-        vibration_profile_creator = VibrationsGraphCreator(self._st_config)
+        gcreator = GraphCreatorFactory.create_graph_creator('vibrations profile', self._st_config)
         st_process = ShakeTuneProcess(
             self._st_config,
             self._printer.get_reactor(),
-            vibration_profile_creator,
+            gcreator,
             self.timeout,
         )
         create_vibrations_profile(gcmd, self._config, st_process)