SIENTIAPDE-1579: Lint fixes and formatting

model_manager/activities/experiment_tracking.py

@@ -6,6 +6,8 @@ in the PostgreSQL database, extending the base Postgres activity with specialize
 methods for experiment management.
 """
 
+import enum
+
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
@@ -13,7 +15,6 @@ with workflow.unsafe.imports_passed_through():
     import traceback
     from collections.abc import Mapping
     from datetime import UTC, datetime
-    from enum import Enum
     from typing import Any
 
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
@@ -24,7 +25,7 @@ with workflow.unsafe.imports_passed_through():
     from sqlalchemy import text
 
 
-class UpdateType(str, Enum):
+class UpdateType(enum.StrEnum):
     """Types of experiment run updates."""
 
     STATUS = 'status'

model_manager/sientia/models.py

@@ -463,7 +463,9 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
         if not date_str:
             return None
         try:
-            python_fmt = _frontend_date_format_to_strftime(self.date_format) if self.date_format else None
+            python_fmt = (
+                _frontend_date_format_to_strftime(self.date_format) if self.date_format else None
+            )
             if python_fmt:
                 return pd.to_datetime(date_str, format=python_fmt)
             return pd.to_datetime(date_str)

model_manager/utils/models/experiment_status.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class ExperimentStatus(str, Enum):
+class ExperimentStatus(StrEnum):
     """
     Status values for experiment run lifecycle.
 

model_manager/utils/repository/training_repository.py

@@ -481,9 +481,9 @@ class TrainingRepository:
             'datetime',
             'DateTime',
         ]
-        timestamp_columns = (
+        timestamp_columns = ([params.date_column] if params.date_column else []) + [
-            [params.date_column] if params.date_column else []
+            c for c in common_timestamp_columns if c != params.date_column
-        ) + [c for c in common_timestamp_columns if c != params.date_column]
+        ]
 
         for col in timestamp_columns:
             if col in data.columns:

scripts/run_cleanup_test.py

@@ -14,18 +14,18 @@ import asyncio
 import os
 import sys
 from datetime import timedelta
-from typing import Any
-from dotenv import load_dotenv
 from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
 from temporalio.client import Client
 
-# Ensure project root is on PYTHONPATH when running directly
+# Ensure project root is on PYTHONPATH when running directly (must run before model_manager import)
 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 if ROOT_DIR not in sys.path:
     sys.path.insert(0, ROOT_DIR)
 
-from model_manager.workflows.cleanup_files import CleanupFiles
+from model_manager.workflows.cleanup_files import CleanupFiles  # noqa: E402
-
 
 # Carrega variáveis de ambiente do arquivo .env na raiz do projeto
 PROJECT_ROOT = Path(__file__).resolve().parent.parent
@@ -53,7 +53,7 @@ async def main(argv: list[str]) -> None:
     if argv:
         bucket_name = argv[0]
 
-    print(f"Connecting to Temporal at {temporal_host} (namespace={temporal_namespace})...")
+    print(f'Connecting to Temporal at {temporal_host} (namespace={temporal_namespace})...')
     client = await Client.connect(
         target_host=temporal_host,
         namespace=temporal_namespace,
@@ -64,12 +64,14 @@ async def main(argv: list[str]) -> None:
         'bucket_name': bucket_name,
     }
 
-    workflow_id = f"cleanup-files-manual-{int(asyncio.get_event_loop().time())}"
+    workflow_id = f'cleanup-files-manual-{int(asyncio.get_event_loop().time())}'
 
-    print(f"Starting cleanup_files workflow once...\n"
+    print(
-          f"  workflow_id = {workflow_id}\n"
+        f'Starting cleanup_files workflow once...\n'
-          f"  task_queue  = {task_queue}\n"
+        f'  workflow_id = {workflow_id}\n'
-          f"  bucket_name = {bucket_name}")
+        f'  task_queue  = {task_queue}\n'
+        f'  bucket_name = {bucket_name}'
+    )
 
     handle = await client.start_workflow(
         CleanupFiles.run,
@@ -79,9 +81,9 @@ async def main(argv: list[str]) -> None:
         run_timeout=timedelta(minutes=10),
     )
 
-    print("Workflow started, waiting for completion...")
+    print('Workflow started, waiting for completion...')
     await handle.result()
-    print("cleanup_files workflow completed successfully.")
+    print('cleanup_files workflow completed successfully.')
 
 
 if __name__ == '__main__':  # pragma: no cover - manual utility script

scripts/run_training_test.py

@@ -39,7 +39,6 @@ from dotenv import load_dotenv
 from psycopg2.extras import Json
 from temporalio import client
 
-
 # Carrega variáveis de ambiente do arquivo .env na raiz do projeto
 PROJECT_ROOT = Path(__file__).resolve().parent.parent
 ENV_PATH = PROJECT_ROOT / '.env'
@@ -325,9 +324,7 @@ def run_local_pipeline(
 
     try:
         train_result = training_repository.train(file_content, train_params)
-        train_result = training_repository.after_train_calculation(
+        train_result = training_repository.after_train_calculation(train_params, train_result)
-            train_params, train_result
-        )
     except Exception as e:
         result['error'] = str(e)
         raise  # re-raise so caller gets full traceback for diagnosis
@@ -448,7 +445,7 @@ def run_single_scenario(scenario_name: str, csv_path: Path) -> dict:
     # Load scenario
     try:
         experiment_request = load_scenario(scenario_name)
-        print(f"  Loaded scenario: {scenario_name}")
+        print(f'  Loaded scenario: {scenario_name}')
     except FileNotFoundError as exc:
         result['error'] = str(exc)
         return result
@@ -459,7 +456,7 @@ def run_single_scenario(scenario_name: str, csv_path: Path) -> dict:
     try:
         uploaded_file_name = upload_to_minio(csv_path)
         result['s3_object_name'] = uploaded_file_name
-        print(f"  Uploaded CSV to MinIO: {uploaded_file_name}")
+        print(f'  Uploaded CSV to MinIO: {uploaded_file_name}')
     except subprocess.CalledProcessError as exc:
         result['error'] = f'Failed to upload file to MinIO: {exc}'
         return result
@@ -471,7 +468,7 @@ def run_single_scenario(scenario_name: str, csv_path: Path) -> dict:
     try:
         experiment_run_id = insert_experiment_run(uploaded_file_name, experiment_request)
         result['experiment_run_id'] = experiment_run_id
-        print(f"  Created experiment_run with ID: {experiment_run_id}")
+        print(f'  Created experiment_run with ID: {experiment_run_id}')
     except psycopg2.Error as exc:
         result['error'] = f'Database error while inserting experiment_run: {exc}'
         return result
@@ -487,7 +484,7 @@ def run_single_scenario(scenario_name: str, csv_path: Path) -> dict:
         workflow_id = asyncio.run(trigger_temporal_workflow(workflow_payload))
         result['workflow_id'] = workflow_id
         result['success'] = True
-        print(f"  Workflow started: {workflow_id}")
+        print(f'  Workflow started: {workflow_id}')
     except Exception as exc:  # noqa: BLE001
         result['error'] = f'Failed to start Temporal workflow: {exc}'
         return result
@@ -507,29 +504,26 @@ def print_summary(results: list[dict]) -> None:
     print('\n' + '=' * 60)
     print('SUMMARY')
     print('=' * 60)
-    print(f"Total: {len(results)} | Passed: {len(passed)} | Failed: {len(failed)}")
+    print(f'Total: {len(results)} | Passed: {len(passed)} | Failed: {len(failed)}')
     print('=' * 60)
 
     if passed:
         print('\n✓ PASSED:')
         for r in passed:
-            print(f"  - {r['scenario']}")
+            print(f'  - {r["scenario"]}')
 
     if failed:
         print('\n✗ FAILED:')
         for r in failed:
-            print(f"  - {r['scenario']}")
+            print(f'  - {r["scenario"]}')
             if r['error']:
-                print(f"    Error: {r['error']}")
+                print(f'    Error: {r["error"]}')
 
     print()
 
 
-def main() -> None:
+def _handle_list_scenarios() -> None:
-    args = parse_args()
+    """Print available scenarios and exit."""
-
-    # List scenarios and exit if requested
-    if args.list:
     scenarios = list_available_scenarios()
     if scenarios:
         print('Available test scenarios:')
@@ -539,8 +533,9 @@ def main() -> None:
         print(f'No scenarios found in {TEST_SCENARIOS_DIR}')
     sys.exit(0)
 
-    # Run all scenarios if requested
+
-    if args.all:
+def _handle_run_all(args: argparse.Namespace) -> None:
+    """Run all scenarios and exit with appropriate code."""
     scenarios = list_available_scenarios()
     if not scenarios:
         print(f'No scenarios found in {TEST_SCENARIOS_DIR}', file=sys.stderr)
@@ -548,7 +543,6 @@ def main() -> None:
 
     print(f'Running {len(scenarios)} scenarios...\n')
     results = []
-
     for i, scenario in enumerate(scenarios, 1):
         print(f'[{i}/{len(scenarios)}] Running scenario: {scenario}')
         result = run_single_scenario(scenario, args.csv)
@@ -557,26 +551,26 @@ def main() -> None:
         print(f'[{i}/{len(scenarios)}] {status} {scenario}\n')
 
     print_summary(results)
-
-        # Exit with error code if any scenario failed
     failed_count = sum(1 for r in results if not r['success'])
     sys.exit(1 if failed_count > 0 else 0)
 
-    # Validate-only mode: no external services
+
-    if args.validate_only:
+def _handle_validate_only(args: argparse.Namespace) -> None:
+    """Validate scenario parameters only and exit."""
     if not args.scenario:
         print('Error: --scenario is required with --validate-only.', file=sys.stderr)
         sys.exit(1)
     result = run_validate_only(args.scenario)
     if result['success']:
-            print(f"Validation OK: {result['scenario']}")
+        print(f'Validation OK: {result["scenario"]}')
     else:
-            print(f"Validation failed: {result['error']}", file=sys.stderr)
+        print(f'Validation failed: {result["error"]}', file=sys.stderr)
         sys.exit(1)
-        return
+    sys.exit(0)
 
-    # Local pipeline: same code path as worker, full traceback on error
+
-    if args.local:
+def _handle_local(args: argparse.Namespace) -> None:
+    """Run local pipeline and exit."""
     if not args.scenario:
         print('Error: --scenario is required with --local.', file=sys.stderr)
         sys.exit(1)
@@ -587,7 +581,7 @@ def main() -> None:
         save_mlflow=args.local_save_mlflow,
     )
     if not result['success']:
-            print(f"Error: {result['error']}", file=sys.stderr)
+        print(f'Error: {result["error"]}', file=sys.stderr)
         sys.exit(1)
     out = {'scenario': result['scenario'], 'success': True}
     if result.get('run_name') is not None:
@@ -595,21 +589,15 @@ def main() -> None:
     if result.get('run_dir') is not None:
         out['run_dir'] = result['run_dir']
     print(json.dumps(out, indent=2))
-        return
 
-    # Require scenario argument if not listing or running all
-    if not args.scenario:
-        print('Error: --scenario or --all is required. Use --list to see available scenarios.', file=sys.stderr)
-        sys.exit(1)
 
-    # Run single scenario (MinIO + Postgres + Temporal)
+def _handle_single_scenario(args: argparse.Namespace) -> None:
+    """Run one scenario (MinIO + Postgres + Temporal) and print result."""
     print(f'Running scenario: {args.scenario}')
     result = run_single_scenario(args.scenario, args.csv)
-
     if not result['success']:
-        print(f"Error: {result['error']}", file=sys.stderr)
+        print(f'Error: {result["error"]}', file=sys.stderr)
         sys.exit(1)
-
     print(
         json.dumps(
             {
@@ -623,5 +611,31 @@ def main() -> None:
     )
 
 
+def main() -> None:
+    args = parse_args()
+
+    if args.list:
+        _handle_list_scenarios()
+
+    if args.all:
+        _handle_run_all(args)
+
+    if args.validate_only:
+        _handle_validate_only(args)
+
+    if args.local:
+        _handle_local(args)
+        return
+
+    if not args.scenario:
+        print(
+            'Error: --scenario or --all is required. Use --list to see available scenarios.',
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    _handle_single_scenario(args)
+
+
 if __name__ == '__main__':
     main()

tests/utils/repository/test_training_repository.py

@@ -991,7 +991,9 @@ class TestConfigureDatetimeIndex:
         assert 'var2' in result.columns
 
     @pytest.mark.filterwarnings('ignore::UserWarning')
-    def test_configure_datetime_index_invalid_timestamp_column(self, training_repo, datetime_params):
+    def test_configure_datetime_index_invalid_timestamp_column(
+        self, training_repo, datetime_params
+    ):
         """Test _configure_datetime_index with invalid timestamp values."""
         data = pd.DataFrame(
             {
@@ -1151,10 +1153,12 @@ class TestApplySupportFilters:
 
     def test_snake_case_upper_lower_line(self):
         """Support filters with upper_line/lower_line (snake_case) filter rows."""
-        data = pd.DataFrame({
+        data = pd.DataFrame(
+            {
                 'x': [1.0, 2.0, 3.0, 4.0],
                 'target': [2.0, 4.0, 6.0, 8.0],
-        })
+            }
+        )
         support_filters = {
             'x': {
                 'upper_line': {'intercept': 1.0, 'angle': 50},
@@ -1167,10 +1171,12 @@ class TestApplySupportFilters:
 
     def test_camel_case_upper_lower_line(self):
         """Support filters with upperLine/lowerLine (camelCase) are accepted."""
-        data = pd.DataFrame({
+        data = pd.DataFrame(
+            {
                 'x': [1.0, 2.0, 3.0],
                 'target': [1.0, 2.0, 3.0],
-        })
+            }
+        )
         support_filters = {
             'x': {
                 'upperLine': {'intercept': 2, 'angle': 5},
@@ -1183,11 +1189,13 @@ class TestApplySupportFilters:
 
     def test_two_variables_ands_masks(self):
         """Two variables apply AND of both masks."""
-        data = pd.DataFrame({
+        data = pd.DataFrame(
+            {
                 'a': [1.0, 2.0, 3.0],
                 'b': [1.0, 2.0, 3.0],
                 'target': [2.0, 2.0, 2.0],
-        })
+            }
+        )
         support_filters = {
             'a': {
                 'upper_line': {'intercept': 10, 'angle': 45},