refactored is_workflow.py to use OOP

src/graph_belts.py

@@ -24,7 +24,6 @@ from common_func import (
     compute_curve_similarity_factor,
     compute_spectrogram,
     detect_peaks,
-    get_git_version,
     parse_log,
     setup_klipper_import,
 )
@@ -453,7 +452,7 @@ def compute_signal_data(data, max_freq):
 ######################################################################
 
 
-def belts_calibration(lognames, klipperdir='~/klipper', max_freq=200.0):
+def belts_calibration(lognames, klipperdir='~/klipper', max_freq=200.0, st_version=None):
     set_locale()
     global shaper_calibrate
     shaper_calibrate = setup_klipper_import(klipperdir)
@@ -530,7 +529,6 @@ def belts_calibration(lognames, klipperdir='~/klipper', max_freq=200.0):
     ax_logo.axis('off')
 
     # Adding Shake&Tune version in the top right corner
-    st_version = get_git_version()
     if st_version is not None:
         fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
 

src/graph_shaper.py

@@ -24,7 +24,6 @@ from common_func import (
     compute_mechanical_parameters,
     compute_spectrogram,
     detect_peaks,
-    get_git_version,
     parse_log,
     setup_klipper_import,
 )
@@ -295,7 +294,7 @@ def plot_spectrogram(ax, t, bins, pdata, peaks, max_freq):
 ######################################################################
 
 
-def shaper_calibration(lognames, klipperdir='~/klipper', max_smoothing=None, scv=5.0, max_freq=200.0):
+def shaper_calibration(lognames, klipperdir='~/klipper', max_smoothing=None, scv=5.0, max_freq=200.0, st_version=None):
     set_locale()
     global shaper_calibrate
     shaper_calibrate = setup_klipper_import(klipperdir)
@@ -387,7 +386,6 @@ def shaper_calibration(lognames, klipperdir='~/klipper', max_smoothing=None, scv
     ax_logo.axis('off')
 
     # Adding Shake&Tune version in the top right corner
-    st_version = get_git_version()
     if st_version is not None:
         fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
 

src/graph_vibrations.py

@@ -24,7 +24,6 @@ matplotlib.use('Agg')
 from common_func import (
     compute_mechanical_parameters,
     detect_peaks,
-    get_git_version,
     identify_low_energy_zones,
     parse_log,
     setup_klipper_import,
@@ -533,7 +532,9 @@ def extract_angle_and_speed(logname):
     return float(angle), float(speed)
 
 
-def vibrations_profile(lognames, klipperdir='~/klipper', kinematics='cartesian', accel=None, max_freq=1000.0):
+def vibrations_profile(
+    lognames, klipperdir='~/klipper', kinematics='cartesian', accel=None, max_freq=1000.0, st_version=None
+):
     set_locale()
     global shaper_calibrate
     shaper_calibrate = setup_klipper_import(klipperdir)
@@ -708,7 +709,6 @@ def vibrations_profile(lognames, klipperdir='~/klipper', kinematics='cartesian',
     ax_logo.axis('off')
 
     # Adding Shake&Tune version in the top right corner
-    st_version = get_git_version()
     if st_version is not None:
         fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
 

src/is_workflow.py

@@ -9,362 +9,390 @@
 #   Use the provided Shake&Tune macros instead!
 
 
+import abc
+import argparse
 import glob
-import optparse
 import os
 import shutil
 import sys
 import tarfile
 import time
 from datetime import datetime
+from pathlib import Path
+
+from git import GitCommandError, Repo
 
 from analyze_axesmap import axesmap_calibration
 from graph_belts import belts_calibration
 from graph_shaper import shaper_calibration
 from graph_vibrations import vibrations_profile
-
+from locale_utils import print_with_c_locale
-RESULTS_FOLDER = os.path.expanduser('~/printer_data/config/K-ShakeTune_results')
-KLIPPER_FOLDER = os.path.expanduser('~/klipper')
-RESULTS_SUBFOLDERS = ['belts', 'inputshaper', 'vibrations']
 
 
-def is_file_open(filepath):
+class Config:
+    KLIPPER_FOLDER = os.path.expanduser('~/klipper')
+    RESULTS_BASE_FOLDER = os.path.expanduser('~/printer_data/config/K-ShakeTune_results')
+    RESULTS_SUBFOLDERS = {'belts': 'belts', 'shaper': 'inputshaper', 'vibrations': 'vibrations'}
+
+    @staticmethod
+    def get_results_folder(type):
+        return os.path.join(Config.RESULTS_BASE_FOLDER, Config.RESULTS_SUBFOLDERS[type])
+
+    @staticmethod
+    def get_git_version():
+        try:
+            # Get the absolute path of the script, resolving any symlinks
+            # Then get 1 times to parent dir to be at the git root folder
+            script_path = Path(__file__).resolve()
+            repo_path = script_path.parents[1]
+            repo = Repo(repo_path)
+            try:
+                version = repo.git.describe('--tags')
+            except GitCommandError:
+                version = repo.head.commit.hexsha[:7]  # If no tag is found, use the simplified commit SHA instead
+            return version
+        except Exception:
+            return None
+
+    @staticmethod
+    def parse_arguments():
+        parser = argparse.ArgumentParser(description='Shake&Tune graphs generation script')
+        parser.add_argument(
+            '-t',
+            '--type',
+            dest='type',
+            choices=['belts', 'shaper', 'vibrations', 'axesmap', 'clean'],
+            required=True,
+            help='Type of output graph to produce',
+        )
+        parser.add_argument(
+            '--accel',
+            type=int,
+            default=None,
+            dest='accel_used',
+            help='Accelerometion used for vibrations profile creation or axes map calibration',
+        )
+        parser.add_argument(
+            '--chip_name',
+            type=str,
+            default='adxl345',
+            dest='chip_name',
+            help='Accelerometer chip name used for vibrations profile creation or axes map calibration',
+        )
+        parser.add_argument(
+            '--max_smoothing',
+            type=float,
+            default=None,
+            dest='max_smoothing',
+            help='Maximum smoothing to allow for input shaper filter recommendations',
+        )
+        parser.add_argument(
+            '--scv',
+            '--square_corner_velocity',
+            type=float,
+            default=5.0,
+            dest='scv',
+            help='Square corner velocity used to compute max accel for input shapers filter recommendations',
+        )
+        parser.add_argument(
+            '-m',
+            '--kinematics',
+            dest='kinematics',
+            default='cartesian',
+            choices=['cartesian', 'corexy'],
+            help='Machine kinematics configuration used for the vibrations profile creation',
+        )
+        parser.add_argument(
+            '-c',
+            '--keep_csv',
+            action='store_true',
+            default=False,
+            dest='keep_csv',
+            help='Whether to keep the raw CSV files after processing in addition to the PNG graphs',
+        )
+        parser.add_argument(
+            '-n',
+            '--keep_results',
+            type=int,
+            default=3,
+            dest='keep_results',
+            help='Number of results to keep in the result folder after each run of the script',
+        )
+        parser.add_argument('--dpi', type=int, default=150, dest='dpi', help='DPI of the output PNG files')
+        parser.add_argument('-v', '--version', action='version', version='%(prog)s ' + Config.get_git_version())
+
+        return parser.parse_args()
+
+
+class FileManager:
+    @staticmethod
+    def wait_file_ready(filepath):
+        file_busy = True
+        loop_count = 0
+        while file_busy:
             for proc in os.listdir('/proc'):
                 if proc.isdigit():
                     for fd in glob.glob(f'/proc/{proc}/fd/*'):
                         try:
                             if os.path.samefile(fd, filepath):
-                        return True
-                except FileNotFoundError:
-                    # Klipper has already released the CSV file
                                 pass
-                except PermissionError:
+                        except FileNotFoundError:  # Klipper has already released the CSV file
-                    # Unable to check for this particular process due to permissions
+                            file_busy = False
+                        except PermissionError:  # Unable to check for this particular process due to permissions
                             pass
-    return False
+            if loop_count > 60:
+                # If Klipper is taking too long to release the file (60 * 1s = 1min), exit the script
+                print_with_c_locale(f'Error: Klipper is taking too long to release {filepath}!')
+                sys.exit(1)
+            else:
+                loop_count += 1
+                time.sleep(1)
+        return
+
+    @staticmethod
+    def ensure_folders_exist():
+        for subfolder in Config.RESULTS_SUBFOLDERS.values():
+            folder = os.path.join(Config.RESULTS_BASE_FOLDER, subfolder)
+            os.makedirs(folder, exist_ok=True)
+
+    @staticmethod
+    def clean_old_files(type, keep_results=3, extension='.png'):
+        folder = Config.get_results_folder(type)
+        files = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(extension)]
+        files.sort(key=os.path.getmtime, reverse=True)
+
+        if 'belts' in folder:
+            if len(files) <= keep_results + 1:
+                return
+            else:  # delete the older files
+                for old_file in files[keep_results + 1 :]:
+                    file_date = '_'.join(os.path.splitext(os.path.basename(old_file))[0].split('_')[1:3])
+                    for suffix in ['A', 'B']:
+                        csv_file = os.path.join(folder, f'belt_{file_date}_{suffix}.csv')
+                        if os.path.exists(csv_file):
+                            os.remove(csv_file)
+                    os.remove(old_file)
+        elif 'shaper' in folder:
+            if len(files) <= 2 * keep_results + 1:
+                return
+            else:  # delete the older files
+                for old_file in files[2 * keep_results + 1 :]:
+                    csv_file = os.path.join(folder, os.path.splitext(os.path.basename(old_file))[0] + '.csv')
+                    if os.path.exists(csv_file):
+                        os.remove(csv_file)
+                    os.remove(old_file)
+        elif 'vibrations' in folder:
+            if len(files) <= keep_results + 1:
+                return
+            else:  # delete the older files
+                for old_file in files[keep_results + 1 :]:
+                    os.remove(old_file)
+                    tar_file = os.path.join(folder, os.path.splitext(os.path.basename(old_file))[0] + '.tar.gz')
+                    if os.path.exists(tar_file):
+                        os.remove(tar_file)
 
 
-def create_belts_graph(keep_csv):
+class GraphCreator(abc.ABC):
-    current_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+    def __init__(self, keep_csv, dpi):
-    lognames = []
+        self._keep_csv = keep_csv
+        self._dpi = dpi
 
+        self._graph_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+        self._version = Config.get_git_version()
+
+    @abc.abstractmethod
+    def create_graph(self):
+        pass
+
+
+class BeltsGraphCreator(GraphCreator):
+    def __init__(self, keep_csv=False, dpi=150):
+        super().__init__(keep_csv, dpi)
+
+        self._type = 'belts'
+        self._folder = Config.get_results_folder(self._type)
+
+    def create_graph(self):
         globbed_files = glob.glob('/tmp/raw_data_axis*.csv')
         if not globbed_files:
-        print('No CSV files found in the /tmp folder to create the belt graphs!')
+            raise FileNotFoundError('no CSV files found in the /tmp folder to create the belt comparison graphs!')
-        sys.exit(1)
         if len(globbed_files) < 2:
-        print('Not enough CSV files found in the /tmp folder. Two files are required for the belt graphs!')
+            raise FileNotFoundError('two CSV files are needed to create the belt comparison graphs!')
-        sys.exit(1)
 
-    sorted_files = sorted(globbed_files, key=os.path.getmtime, reverse=True)
+        lognames = []
-
+        for filename in sorted(globbed_files, key=os.path.getmtime, reverse=True)[:2]:
-    for filename in sorted_files[:2]:
             # Wait for the file handler to be released by Klipper
-        while is_file_open(filename):
+            FileManager.wait_file_ready(filename)
-            time.sleep(2)
 
-        # Extract the tested belt from the filename and rename/move the CSV file to the result folder
+            # Cleanup of the filename and moving it in the result folder
             belt = os.path.basename(filename).split('_')[3].split('.')[0].upper()
-        new_file = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[0], f'belt_{current_date}_{belt}.csv')
+            new_file = os.path.join(self._folder, f'{self._type}_{self._graph_date}_{belt}.csv')
             shutil.move(filename, new_file)
-        os.sync()  # Sync filesystem to avoid problems
-
-        # Save the file path for later
             lognames.append(new_file)
 
-        # Wait for the file handler to be released by the move command
+            # Check if the file is ready to be read
-        while is_file_open(new_file):
+            os.sync()
-            time.sleep(2)
+            FileManager.wait_file_ready(new_file)
 
-    # Generate the belts graph and its name
+        fig = belts_calibration(lognames, Config.KLIPPER_FOLDER, self._version)
-    fig = belts_calibration(lognames, KLIPPER_FOLDER)
+        png_filename = os.path.join(self._folder, f'{self._type}_{self._graph_date}.png')
-    png_filename = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[0], f'belts_{current_date}.png')
+        fig.savefig(png_filename, dpi=self._dpi)
-    fig.savefig(png_filename, dpi=150)
 
         # Remove the CSV files if the user don't want to keep them
-    if not keep_csv:
+        if not self._keep_csv:
             for csv in lognames:
                 if os.path.exists(csv):
                     os.remove(csv)
 
-    return
 
+class ShaperGraphCreator(GraphCreator):
+    def __init__(self, max_smoothing=None, scv=5.0, keep_csv=False, dpi=150):
+        super().__init__(keep_csv, dpi)
 
-def create_shaper_graph(keep_csv, max_smoothing, scv):
+        self._max_smoothing = max_smoothing
-    current_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+        self._scv = scv
 
-    # Get all the files and sort them based on last modified time to select the most recent one
+        self._type = 'shaper'
+        self._folder = Config.get_results_folder(self._type)
+
+    def create_graph(self):
         globbed_files = glob.glob('/tmp/raw_data*.csv')
         if not globbed_files:
-        print('No CSV files found in the /tmp folder to create the input shaper graphs!')
+            raise FileNotFoundError('no CSV files found in the /tmp folder to create the input shaper graphs!')
-        sys.exit(1)
 
-    sorted_files = sorted(globbed_files, key=os.path.getmtime, reverse=True)
+        # Find the CSV files with the latest timestamp and wait for it to be released by Klipper
-    filename = sorted_files[0]
+        filename = sorted(globbed_files, key=os.path.getmtime, reverse=True)[0]
+        FileManager.wait_file_ready(filename)
 
-    # Wait for the file handler to be released by Klipper
+        # Cleanup of the filename and moving it in the result folder
-    while is_file_open(filename):
-        time.sleep(2)
-
-    # Extract the tested axis from the filename and rename/move the CSV file to the result folder
         axis = os.path.basename(filename).split('_')[3].split('.')[0].upper()
-    new_file = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[1], f'resonances_{current_date}_{axis}.csv')
+        new_file = os.path.join(self._folder, f'{self._type}_{self._graph_date}_{axis}.csv')
         shutil.move(filename, new_file)
-    os.sync()  # Sync filesystem to avoid problems
 
-    # Wait for the file handler to be released by the move command
+        # Check if the file is ready to be read
-    while is_file_open(new_file):
+        os.sync()
-        time.sleep(2)
+        FileManager.wait_file_ready(new_file)
 
-    # Generate the shaper graph and its name
+        fig = shaper_calibration(
-    fig = shaper_calibration([new_file], KLIPPER_FOLDER, max_smoothing=max_smoothing, scv=scv)
+            [new_file],
-    png_filename = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[1], f'resonances_{current_date}_{axis}.png')
+            Config.KLIPPER_FOLDER,
-    fig.savefig(png_filename, dpi=150)
+            max_smoothing=self._max_smoothing,
+            scv=self._scv,
+            st_version=self._version,
+        )
+        png_filename = os.path.join(self._folder, f'{self._type}_{self._graph_date}_{axis}.png')
+        fig.savefig(png_filename, dpi=self._dpi)
 
-    # Remove the CSV file if the user don't want to keep it
+        # Remove the CSV files if the user don't want to keep them
-    if not keep_csv:
+        if not self._keep_csv:
             if os.path.exists(new_file):
                 os.remove(new_file)
 
         return axis
 
 
-def create_vibrations_graph(accel, kinematics, chip_name, keep_csv):
+class VibrationsGraphCreator(GraphCreator):
-    current_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+    def __init__(self, kinematics, accel, chip_name, keep_csv=False, dpi=150):
-    lognames = []
+        super().__init__(keep_csv, dpi)
 
-    globbed_files = glob.glob(f'/tmp/{chip_name}-*.csv')
+        self._kinematics = kinematics
+        self._accel = accel
+        self._chip_name = chip_name
+
+        self._type = 'vibrations'
+        self._folder = Config.get_results_folder(self._type)
+
+    def create_graph(self):
+        globbed_files = glob.glob(f'/tmp/{self._chip_name}-*.csv')
         if not globbed_files:
-        print('No CSV files found in the /tmp folder to create the vibration graphs!')
+            raise FileNotFoundError('no CSV files found in the /tmp folder to create the vibrations graphs!')
-        sys.exit(1)
         if len(globbed_files) < 3:
-        print('Not enough CSV files found in the /tmp folder. At least 3 files are required for the vibration graphs!')
+            raise FileNotFoundError('at least 3 CSV files are needed to create the vibrations graphs!')
-        sys.exit(1)
 
+        lognames = []
         for filename in globbed_files:
             # Wait for the file handler to be released by Klipper
-        while is_file_open(filename):
+            FileManager.wait_file_ready(filename)
-            time.sleep(2)
 
             # Cleanup of the filename and moving it in the result folder
-        cleanfilename = os.path.basename(filename).replace(chip_name, f'vibr_{current_date}')
+            cleanfilename = os.path.basename(filename).replace(self._chip_name, f'{self._type}_{self._graph_date}')
-        new_file = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[2], cleanfilename)
+            new_file = os.path.join(self._folder, cleanfilename)
             shutil.move(filename, new_file)
-
-        # Save the file path for later
             lognames.append(new_file)
 
         # Sync filesystem to avoid problems as there is a lot of file copied
         os.sync()
         time.sleep(5)
 
-    # Generate the vibration graph and its name
+        fig = vibrations_profile(lognames, Config.KLIPPER_FOLDER, self._kinematics, self._accel, self._version)
-    fig = vibrations_profile(lognames, KLIPPER_FOLDER, kinematics, accel)
+        png_filename = os.path.join(self._folder, f'{self._type}_{self._graph_date}.png')
-    png_filename = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[2], f'vibrations_{current_date}.png')
+        fig.savefig(png_filename, dpi=self._dpi)
-    fig.savefig(png_filename, dpi=150)
 
         # Archive all the csv files in a tarball in case the user want to keep them
-    if keep_csv:
+        if self._keep_csv:
-        with tarfile.open(
+            with tarfile.open(os.path.join(self._folder, f'{self._type}_{self._graph_date}.tar.gz'), 'w:gz') as tar:
-            os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[2], f'vibrations_{current_date}.tar.gz'), 'w:gz'
-        ) as tar:
                 for csv_file in lognames:
                     tar.add(csv_file, arcname=os.path.basename(csv_file), recursive=False)
 
         # Remove the remaining CSV files not needed anymore (tarball is safe if it was created)
-    for csv_file in lognames:
+        for csv in lognames:
-        if os.path.exists(csv_file):
+            if os.path.exists(csv):
-            os.remove(csv_file)
+                os.remove(csv)
-
-    return
 
 
-def find_axesmap(accel, chip_name):
+class AxesMapFinder:
-    current_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+    def __init__(self, accel, chip_name):
-    result_filename = os.path.join(RESULTS_FOLDER, f'axes_map_{current_date}.txt')
+        self._accel = accel
-    lognames = []
+        self._chip_name = chip_name
 
-    globbed_files = glob.glob(f'/tmp/{chip_name}-*.csv')
+        self._graph_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+
+        self._type = 'axesmap'
+        self._folder = Config.RESULTS_BASE_FOLDER
+
+    def find_axesmap(self):
+        globbed_files = glob.glob(f'/tmp/{self._chip_name}-*.csv')
         if not globbed_files:
-        print('No CSV files found in the /tmp folder to analyze and find the axes_map!')
+            raise FileNotFoundError('no CSV files found in the /tmp folder to find the axes map!')
-        sys.exit(1)
 
-    sorted_files = sorted(globbed_files, key=os.path.getmtime, reverse=True)
+        # Find the CSV files with the latest timestamp and wait for it to be released by Klipper
-    filename = sorted_files[0]
+        logname = sorted(globbed_files, key=os.path.getmtime, reverse=True)[0]
-
+        FileManager.wait_file_ready(logname)
-    # Wait for the file handler to be released by Klipper
-    while is_file_open(filename):
-        time.sleep(2)
-
-    # Analyze the CSV to find the axes_map parameter
-    lognames.append(filename)
-    results = axesmap_calibration(lognames, accel)
 
+        results = axesmap_calibration([logname], self._accel)
+        result_filename = os.path.join(self._folder, f'{self._type}_{self._graph_date}.txt')
         with open(result_filename, 'w') as f:
             f.write(results)
 
-    return
-
-
-# Utility function to get old files based on their modification time
-def get_old_files(folder, extension, limit):
-    files = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(extension)]
-    files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
-    return files[limit:]
-
-
-def clean_files(keep_results):
-    # Define limits based on STORE_RESULTS
-    keep1 = keep_results + 1
-    keep2 = 2 * keep_results + 1
-
-    # Find old files in each directory
-    old_belts_files = get_old_files(os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[0]), '.png', keep1)
-    old_inputshaper_files = get_old_files(os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[1]), '.png', keep2)
-    old_speed_vibr_files = get_old_files(os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[2]), '.png', keep1)
-
-    # Remove the old belt files
-    for old_file in old_belts_files:
-        file_date = '_'.join(os.path.splitext(os.path.basename(old_file))[0].split('_')[1:3])
-        for suffix in ['A', 'B']:
-            csv_file = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[0], f'belt_{file_date}_{suffix}.csv')
-            if os.path.exists(csv_file):
-                os.remove(csv_file)
-        os.remove(old_file)
-
-    # Remove the old shaper files
-    for old_file in old_inputshaper_files:
-        csv_file = os.path.join(
-            RESULTS_FOLDER, RESULTS_SUBFOLDERS[1], os.path.splitext(os.path.basename(old_file))[0] + '.csv'
-        )
-        if os.path.exists(csv_file):
-            os.remove(csv_file)
-        os.remove(old_file)
-
-    # Remove the old vibrations files
-    for old_file in old_speed_vibr_files:
-        os.remove(old_file)
-        tar_file = os.path.join(
-            RESULTS_FOLDER, RESULTS_SUBFOLDERS[2], os.path.splitext(os.path.basename(old_file))[0] + '.tar.gz'
-        )
-        if os.path.exists(tar_file):
-            os.remove(tar_file)
-
 
 def main():
-    # Parse command-line arguments
+    print_with_c_locale(f'Shake&Tune version: {Config.get_git_version()}')
-    usage = '%prog [options] <logs>'
-    opts = optparse.OptionParser(usage)
-    opts.add_option('-t', '--type', type='string', dest='type', default=None, help='type of output graph to produce')
-    opts.add_option(
-        '--accel',
-        type='int',
-        default=None,
-        dest='accel_used',
-        help='acceleration used during the vibration macro or axesmap macro',
-    )
-    opts.add_option(
-        '--axis_name', type='string', default=None, dest='axis_name', help='axis tested during the vibration macro'
-    )
-    opts.add_option(
-        '--chip_name',
-        type='string',
-        default='adxl345',
-        dest='chip_name',
-        help='accelerometer chip name in klipper used during the vibration macro or the axesmap macro',
-    )
-    opts.add_option(
-        '-n',
-        '--keep_results',
-        type='int',
-        default=3,
-        dest='keep_results',
-        help='number of results to keep in the result folder after each run of the script',
-    )
-    opts.add_option(
-        '-c',
-        '--keep_csv',
-        action='store_true',
-        default=False,
-        dest='keep_csv',
-        help='weither or not to keep the CSV files alongside the PNG graphs image results',
-    )
-    opts.add_option(
-        '--scv',
-        '--square_corner_velocity',
-        type='float',
-        dest='scv',
-        default=5.0,
-        help='square corner velocity used to compute max accel for axis shapers graphs',
-    )
-    opts.add_option(
-        '--max_smoothing', type='float', dest='max_smoothing', default=None, help='maximum shaper smoothing to allow'
-    )
-    opts.add_option(
-        '-m',
-        '--kinematics',
-        type='string',
-        dest='kinematics',
-        default='cartesian',
-        help='machine kinematics configuration used for the vibrations graphs',
-    )
-    options, _ = opts.parse_args()
 
-    if options.type is None:
+    options = Config.parse_arguments()
-        opts.error('You must specify the type of output graph you want to produce (option -t)')
+    FileManager.ensure_folders_exist()
-    elif options.type.lower() is None or options.type.lower() not in [
+
-        'belts',
+    graph_creator = None
-        'shaper',
+    if options.type == 'belts':
-        'vibrations',
+        graph_creator = BeltsGraphCreator(options.keep_csv, options.dpi)
-        'axesmap',
+    elif options.type == 'shaper':
-        'clean',
+        graph_creator = ShaperGraphCreator(options.max_smoothing, options.scv, options.keep_csv, options.dpi)
-    ]:
+    elif options.type == 'vibrations':
-        opts.error(
+        graph_creator = VibrationsGraphCreator(
-            "Type of output graph need to be in the list of 'belts', 'shaper', 'vibrations', 'axesmap' or 'clean'"
+            options.kinematics, options.accel_used, options.chip_name, options.keep_csv, options.dpi
         )
-    else:
+    elif options.type == 'axesmap':
-        graph_mode = options.type
+        graph_creator = AxesMapFinder(options.accel_used, options.chip_name)
+    elif options.type == 'clean':
+        print_with_c_locale(f'Cleaning output folder to keep only the last {options.keep_results} results...')
+        FileManager.clean_old_files(options.type, options.keep_results)
+        return
 
-    if graph_mode.lower() == 'vibrations' and options.kinematics not in ['cartesian', 'corexy']:
+    if graph_creator:
-        opts.error('Only Cartesian and CoreXY kinematics are supported by this tool at the moment!')
+        graph_creator.create_graph()
+        print_with_c_locale(f'{options.type} graphs created successfully!')
 
-    # Check if results folders are there or create them before doing anything else
+    if options.keep_csv is False and options.type != 'clean':
-    for result_subfolder in RESULTS_SUBFOLDERS:
+        print_with_c_locale('Deleting raw CSV files... If you want to keep them, use the --keep_csv option!')
-        folder = os.path.join(RESULTS_FOLDER, result_subfolder)
-        if not os.path.exists(folder):
-            os.makedirs(folder)
-
-    if graph_mode.lower() == 'belts':
-        create_belts_graph(keep_csv=options.keep_csv)
-        print(f'Belt graph created. You will find the results in {RESULTS_FOLDER}/{RESULTS_SUBFOLDERS[0]}')
-
-    elif graph_mode.lower() == 'shaper':
-        axis = create_shaper_graph(keep_csv=options.keep_csv, max_smoothing=options.max_smoothing, scv=options.scv)
-        print(
-            f'{axis} input shaper graph created. You will find the results in {RESULTS_FOLDER}/{RESULTS_SUBFOLDERS[1]}'
-        )
-
-    elif graph_mode.lower() == 'vibrations':
-        create_vibrations_graph(
-            accel=options.accel_used,
-            kinematics=options.kinematics,
-            chip_name=options.chip_name,
-            keep_csv=options.keep_csv,
-        )
-        print(f'Vibrations graph created. You will find the results in {RESULTS_FOLDER}/{RESULTS_SUBFOLDERS[2]}')
-
-    elif graph_mode.lower() == 'axesmap':
-        print(
-            'WARNING: AXES_MAP_CALIBRATION is currently very experimental and may produce incorrect results... Please validate the output!'
-        )
-        find_axesmap(accel=options.accel_used, chip_name=options.chip_name)
-
-    elif graph_mode.lower() == 'clean':
-        print(f'Cleaning output folder to keep only the last {options.keep_results} results...')
-        clean_files(keep_results=options.keep_results)
-
-    if options.keep_csv is False and graph_mode.lower() != 'clean':
-        print('Deleting raw CSV files... If you want to keep them, use the --keep_csv option!')
 
 
 if __name__ == '__main__':

src/locale_utils.py

@@ -6,6 +6,7 @@
 
 import locale
 
+
 # Set the best locale for time and date formating (generation of the titles)
 def set_locale():
     try:
@@ -15,16 +16,19 @@ def set_locale():
     except locale.Error:
         locale.setlocale(locale.LC_TIME, 'C')
 
+
 # Print function to avoid problem in Klipper console (that doesn't support special characters) due to locale settings
 def print_with_c_locale(*args, **kwargs):
     try:
         original_locale = locale.getlocale()
         locale.setlocale(locale.LC_ALL, 'C')
     except locale.Error as e:
-        print("Warning: Failed to set a basic locale. Special characters may not display correctly in Klipper console:", e)
+        print(
+            'Warning: Failed to set a basic locale. Special characters may not display correctly in Klipper console:', e
+        )
     finally:
         print(*args, **kwargs)  # Proceed with printing regardless of locale setting success
         try:
             locale.setlocale(locale.LC_ALL, original_locale)
         except locale.Error as e:
-            print("Warning: Failed to restore the original locale setting:", e)
+            print('Warning: Failed to restore the original locale setting:', e)