Front-end Improvements: main.py refactoring, ConfigSettings enhan…

…cement

* Improve `ConfigSettings` class, add class method to write out default config, update tests
* Refactor the entry script `scripts/main.py`, add cmdline option to write out default yaml

point_utils/settings.py

@@ -5,34 +5,61 @@
 from pathlib import Path
 from pydantic import BaseModel, Field, field_validator
 
+
 class ConfigSettings(BaseModel):
     """
     Schema for the configuration settings.
     """
     # Required settings
     input_file: str = Field(description="Path to the input text file.")
     offset_magnitude: float = Field(description="Magnitude of the offset.")
+    data_label_to_offset: float = Field(
+        description="Label of the points to offset.")
+    new_data_label: float = Field(
+        description="Label for the new offset points.")
 
     # Optional settings
     offset_method: str = Field(
-        default="KDTreeOffsets", description="Method used for offset calculation.")
-    output_file: str = Field(
-        default="output.txt", description="Path to the output text file.")
-    visualize: bool = Field(
-        default=False, description="Plot the final output point cloud.")
+        default="KDTreeOffsets",
+        description="Method used for offset calculation.")
+    output_file: str = Field(default="output.txt",
+                             description="Path to the output text file.")
+    visualize: bool = Field(default=False,
+                            description="Plot the final output point cloud.")
 
     @field_validator("input_file")
     def check_input_file(cls, input_file):
         if Path(input_file).suffix != ".txt":
             raise ValueError("Input file must be a text file.")
         return input_file
 
+    @classmethod
+    def write_default_config_to_yaml(cls, yaml_path: str):
+        """
+        Write the default config YAML file.
+        """
+        with open(yaml_path, "w") as file:
+            file.write(cls._get_default_settings_yaml_string())
+
+    @classmethod
+    def _get_default_settings_yaml_string(cls):
+        """
+        Get the default settings to a YAML string.
+        """
+        schema = cls.schema()
+        default_strings = []
+        for key, val_dict in schema["properties"].items():
+            default_value = val_dict.get("default", ' ')
+            default_strings.append(f"{key}: {default_value}")
+
+        return "\n".join(default_strings)
+
 
-def load_and_validate_config(yaml_file_path: str) -> ConfigSettings:
+def load_and_validate_config(yaml_path: str) -> ConfigSettings:
     """
     Read the YAML configuration file and validate the settings.
     """
-    with open(yaml_file_path, "r") as file:
+    with open(yaml_path, "r") as file:
         yaml_data = yaml.safe_load(file)
 
     # Validate and parse using the Pydantic model

scripts/main.py

@@ -2,25 +2,34 @@
 import pprint
 import argparse
 from pathlib import Path
-import numpy as np
 
 from point_utils import DEFAULT_OUT_TXT
+from point_utils import ConfigSettings
 from point_utils.utils import get_data_from_txt, save_to_txt, visualize
 from point_utils.offsetter import offset_factory
 
-LABLE_B = "B"
-LABLE_C = "C"
 
+def _parse_args():
+    """
+    Define command-line interface
+    """
+    parser = argparse.ArgumentParser(
+        description="Run offset points calculation")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("-config",
+                       type=str,
+                       help="Path to the YAML configuration file")
+    group.add_argument('-write-default-config',
+                       type=str,
+                       help="Write the default configuration file")
+    return parser.parse_args()
 
-def main():
-    parser = argparse.ArgumentParser(description="Run offset point calculation")
-    parser.add_argument(
-        "-config", type=str, required=True, help="Path to the YAML configuration file"
-    )
-    args = parser.parse_args()
-
-    config_path = args.config
-    assert Path(config_path).exists(), f"Configuration file not found: {config_path}"
+
+def run(config_path: str):
+    if not Path(config_path).exists():
+        raise FileNotFoundError(f"Configuration file not found: {config_path}")
+
+    # Load settings from the configuration file
     print(f"Using configuration file: {config_path}")
     with open(config_path, 'r') as fhandle:
         settings = yaml.safe_load(fhandle)
@@ -29,28 +38,40 @@ def main():
 
     all_coordinates, labels = get_data_from_txt(settings['input_file'])
 
-    # Calculate offset points for selected points
-    offset_calculator = offset_factory('KDTreeOffsets',
-                                       all_coordinates=all_coordinates,
-                                       labels=labels,
-                                       # LABLE_B is used to indicate points to offset.
-                                       data_label_to_offset=LABLE_B,
-                                       offset_magnitude=settings['offset_magnitude'],
-                                       new_data_label=LABLE_C)
-    # Set method-specific keyword arguments
+    # Setup the offset calculator based on settings
+    offset_calculator = offset_factory(
+        offset_method=settings['offset_method'],
+        all_coordinates=all_coordinates,
+        labels=labels,
+        data_label_to_offset=settings['data_label_to_offset'],
+        offset_magnitude=settings['offset_magnitude'],
+        new_data_label=settings['new_data_label'])
+
+    # Set method-specific keyword arguments (if any)
     kwargs = {}
+
+    # Calculate offset points for selected points
     offset_calculator.add_offset_points(**kwargs)
 
     out_file = Path(settings.get('output_file', DEFAULT_OUT_TXT))
     out_path = out_file.resolve()
-    save_to_txt(out_path, offset_calculator.all_coordinates, offset_calculator.labels)
-    print(f"Save data to {out_path}")
+    save_to_txt(out_path, offset_calculator.all_coordinates,
+                offset_calculator.labels)
+    print(f"Save data to {out_path}.")
 
     if settings.get('visualize', True):
-        print(f"Visualizing the data in {out_path}")
-        visualize(out_file)
+        print("Visualizing the data.")
+        visualize(out_file, output_file=out_file.with_suffix(".png"))
+
+
+def main():
+    args = _parse_args()
+    if args.write_default_config:
+        ConfigSettings.write_default_config_to_yaml(args.write_default_config)
+        return
+
+    run(args.config)
 
 
 if __name__ == "__main__":
     main()
-

tests/settings_test.py

@@ -17,6 +17,9 @@ def test_imported():
 
 @pytest.fixture(autouse=True)
 def tmp_cwd():
+    """
+    Change to a temporary directory for the duration of the test
+    """
     old_cwd = os.getcwd()
     tmp_dir = tempfile.mkdtemp()
     os.chdir(tmp_dir)
@@ -26,18 +29,21 @@ def tmp_cwd():
 
 def test_config_missing_required_setting():
     with pytest.raises(
-        ValidationError,
-        match="2 validation errors for ConfigSettings") as exc_info:
+            ValidationError,
+            match="4 validation errors for ConfigSettings") as exc_info:
         config = ConfigSettings()
     assert "input_file" in str(exc_info.value)
     assert "offset_magnitude" in str(exc_info.value)
 
 
 def test_config_invalid_input_ext(tmp_cwd):
-    setting = {"input_file": "data.json", "offset_magnitude": 2.0}
+    setting = {
+        "input_file": "data.json",
+        "offset_magnitude": 2.0,
+        "data_label_to_offset": "B",
+        "new_data_label": "C",
+    }
     Path("config.yaml").write_text(yaml.dump(setting))
-    with pytest.raises(
-        ValidationError,
-        match="Value error, Input file must be a text file."):
+    with pytest.raises(ValidationError,
+                       match="Value error, Input file must be a text file."):
         config = load_and_validate_config("config.yaml")
-