Refactor orchestrator module to add Literal type hinting for config_params argument of run_pipeline; update imports across multiple modules for consistency of importing only at module-level.

Author: alperkent-cmi

src/graphomotor/core/config.py

@@ -1,13 +1,13 @@
 """Configuration module for Graphomotor."""
 
+import dataclasses
 import logging
 import warnings
-from dataclasses import dataclass
 
 import numpy as np
 
 
-@dataclass
+@dataclasses.dataclass
 class SpiralConfig:
     """Class for the parameters of anticipated spiral drawing."""
 

src/graphomotor/core/models.py

@@ -1,14 +1,14 @@
 """Internal data class for spiral drawing data."""
 
-from datetime import datetime
-from typing import Callable
+import datetime
+import typing
 
 import numpy as np
 import pandas as pd
-from pydantic import BaseModel, ConfigDict, field_validator
+import pydantic
 
 
-class Spiral(BaseModel):
+class Spiral(pydantic.BaseModel):
     """Class representing a spiral drawing, encapsulating both raw data and metadata.
 
     Attributes:
@@ -21,12 +21,12 @@ class Spiral(BaseModel):
             - start_time: Start time of drawing.
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
     data: pd.DataFrame
-    metadata: dict[str, str | datetime]
+    metadata: dict[str, str | datetime.datetime]
 
-    @field_validator("data")
+    @pydantic.field_validator("data")
     @classmethod
     def validate_dataframe(cls, v: pd.DataFrame) -> pd.DataFrame:
         """Validate that DataFrame contains required columns and correct data types.
@@ -46,7 +46,7 @@ def validate_dataframe(cls, v: pd.DataFrame) -> pd.DataFrame:
 
         return v
 
-    @field_validator("metadata")
+    @pydantic.field_validator("metadata")
     @classmethod
     def validate_metadata(cls, v: dict) -> dict:
         """Validate metadata dictionary for required keys and correct data types.
@@ -102,7 +102,7 @@ def all(cls) -> set[str]:
     @classmethod
     def get_extractors(
         cls, spiral: Spiral, reference_spiral: np.ndarray
-    ) -> dict[str, Callable[[], dict[str, float]]]:
+    ) -> dict[str, typing.Callable[[], dict[str, float]]]:
         """Get all feature extractors with appropriate inputs.
 
         Args:

src/graphomotor/core/orchestrator.py

@@ -1,9 +1,9 @@
 """Runner for the Graphomotor pipeline."""
 
+import datetime
 import os
 import pathlib
-from datetime import datetime
-from typing import Literal
+import typing
 
 import numpy as np
 import pandas as pd
@@ -14,7 +14,7 @@
 
 logger = config.get_logger()
 
-FeatureCategories = Literal["duration", "velocity", "hausdorff", "AUC"]
+FeatureCategories = typing.Literal["duration", "velocity", "hausdorff", "AUC"]
 
 
 def _ensure_path(path: pathlib.Path | str) -> pathlib.Path:
@@ -122,7 +122,7 @@ def _export_features_to_csv(
 
     filename = (
         f"{participant_id}_{task}_{hand}_features_"
-        f"{datetime.today().strftime('%Y%m%d')}.csv"
+        f"{datetime.datetime.today().strftime('%Y%m%d')}.csv"
     )
 
     if not output_path.suffix:
@@ -220,7 +220,19 @@ def run_pipeline(
         "hausdorff",
         "AUC",
     ],
-    config_params: dict[str, float | int] | None = None,
+    config_params: dict[
+        typing.Literal[
+            "center_x",
+            "center_y",
+            "start_radius",
+            "growth_rate",
+            "start_angle",
+            "end_angle",
+            "num_points",
+        ],
+        float | int,
+    ]
+    | None = None,
 ) -> dict[str, str]:
     """Run the Graphomotor pipeline to extract features from spiral drawings.
 
@@ -238,14 +250,14 @@ def run_pipeline(
             - "AUC": Area under the curve metric
         config_params: Optional dictionary with custom spiral configuration parameters.
             These parameters control reference spiral generation and spiral centering.
-            If None, default parameters are used. Supported parameters are:
-            - "center_x": X-coordinate of the spiral center. Default is 50.
-            - "center_y": Y-coordinate of the spiral center. Default is 50.
-            - "start_radius": Starting radius of the spiral. Default is 0.
-            - "growth_rate": Growth rate of the spiral. Default is 1.075.
-            - "start_angle": Starting angle of the spiral. Default is 0.
-            - "end_angle": Ending angle of the spiral. Default is 8π.
-            - "num_points": Number of points in the spiral. Default is 10000.
+            If None, default configuration is used. Supported parameters are:
+            - "center_x" (float): X-coordinate of the spiral center. Default is 50.
+            - "center_y" (float): Y-coordinate of the spiral center. Default is 50.
+            - "start_radius" (float): Starting radius of the spiral. Default is 0.
+            - "growth_rate" (float): Growth rate of the spiral. Default is 1.075.
+            - "start_angle" (float): Starting angle of the spiral. Default is 0.
+            - "end_angle" (float): Ending angle of the spiral. Default is 8π.
+            - "num_points" (int): Number of points in the spiral. Default is 10000.
 
     Returns:
         Dictionary of extracted features.
@@ -258,7 +270,9 @@ def run_pipeline(
     spiral_config = None
     if config_params:
         logger.info(f"Custom spiral configuration: {config_params}")
-        spiral_config = config.SpiralConfig.add_custom_params(config_params)
+        spiral_config = config.SpiralConfig.add_custom_params(
+            typing.cast(dict, config_params)
+        )
 
     features = extract_features(
         input_path, output_path, feature_categories, spiral_config

src/graphomotor/utils/center_spiral.py

@@ -1,15 +1,15 @@
 """Utility functions for centering a spiral."""
 
-from typing import overload
+import typing
 
 import numpy as np
 
 from graphomotor.core import config, models
 
 
-@overload
+@typing.overload
 def center_spiral(spiral: models.Spiral) -> models.Spiral: ...
-@overload
+@typing.overload
 def center_spiral(spiral: np.ndarray) -> np.ndarray: ...
 def center_spiral(spiral):
     """Center a spiral by translating it to the origin.

tests/unit/test_center_spiral.py

@@ -1,6 +1,6 @@
 """Test cases for center_spiral.py functions."""
 
-from typing import cast
+import typing
 
 import numpy as np
 import pytest
@@ -42,4 +42,4 @@ def test_center_spiral_invalid_type(
         TypeError,
         match=f"Expected models.Spiral or np.ndarray, got {expected_type_name}",
     ):
-        center_spiral.center_spiral(cast(models.Spiral, invalid_input))
+        center_spiral.center_spiral(typing.cast(models.Spiral, invalid_input))

tests/unit/test_orchestrator.py

@@ -1,8 +1,8 @@
 """Tests for the orchestrator module."""
 
+import datetime
 import pathlib
-from datetime import datetime
-from typing import cast
+import typing
 
 import numpy as np
 import pytest
@@ -48,7 +48,7 @@ def test_validate_feature_categories_only_invalid(
 
 def test_validate_feature_categories_mixed(caplog: pytest.LogCaptureFixture) -> None:
     """Test _validate_feature_categories with mix of valid and invalid categories."""
-    feature_categories = cast(
+    feature_categories = typing.cast(
         list[orchestrator.FeatureCategories],
         [
             "duration",
@@ -136,7 +136,7 @@ def test_export_features_to_csv_no_extension_dir_exists(
 
     expected_filename = (
         f"{valid_spiral.metadata['id']}_{valid_spiral.metadata['task']}_{valid_spiral.metadata['hand']}_"
-        f"features_{datetime.today().strftime('%Y%m%d')}.csv"
+        f"features_{datetime.datetime.today().strftime('%Y%m%d')}.csv"
     )
 
     expected_output_path = output_path / expected_filename
@@ -162,7 +162,7 @@ def test_export_features_to_csv_no_extension_no_dir(
 
     expected_filename = (
         f"{valid_spiral.metadata['id']}_{valid_spiral.metadata['task']}_{valid_spiral.metadata['hand']}_"
-        f"features_{datetime.today().strftime('%Y%m%d')}.csv"
+        f"features_{datetime.datetime.today().strftime('%Y%m%d')}.csv"
     )
 
     expected_output_path = output_path / expected_filename