Add progress callback for clean pipeline (#114)

* Add progress callback for clean pipeline

* Align progress callback event schema

* Update src/freshdata/cleaner.py

Co-authored-by: Rohith <mamidalarohith7@gmail.com>

---------

Co-authored-by: Johnny Wilson Dougherty <johnnydougherty09@gmail.com>
Co-authored-by: Rohith <mamidalarohith7@gmail.com>

Author: 3 people

src/freshdata/api.py

@@ -226,8 +226,9 @@ def clean(
         ``strategy`` (``"balanced"`` default / ``"aggressive"`` / ``"conservative"``),
         ``missing_threshold_low``/``_medium``/``_high``, ``duplicate_threshold``,
         ``outlier_method``, ``outlier_action``, ``preserve_original``, ``verbose``,
-        ``preserve_columns``, ``target_column``, ``duplicate_keep``, ``impute``,
-        ``outliers``. Unknown names raise :class:`TypeError`.
+        ``progress_callback``, ``preserve_columns``, ``target_column``,
+        ``duplicate_keep``, ``impute``, ``outliers``. Unknown names raise
+        :class:`TypeError`.
 
     Examples
     --------

src/freshdata/cleaner.py

@@ -4,7 +4,7 @@
 
 import dataclasses
 import time
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 
 import pandas as pd
 
@@ -23,6 +23,8 @@
 from .steps.prune import drop_constant_columns, drop_empty_columns, drop_empty_rows
 from .steps.strings import clean_strings
 
+ProgressCallback = Callable[[dict[str, object]], None]
+
 
 def _validate_input(df: object, config: CleanConfig) -> pd.DataFrame:
     if isinstance(df, pd.Series):
@@ -42,9 +44,27 @@ def _validate_input(df: object, config: CleanConfig) -> pd.DataFrame:
     return frame
 
 
+def _emit_progress(
+    callback: ProgressCallback | None,
+    step: str,
+    status: str,
+    frame: pd.DataFrame,
+) -> None:
+    if callback is None:
+        return
+    callback(
+        {
+            "step": step,
+            "status": status,
+            "rows": len(frame),
+            "columns": frame.shape[1],
+        }
+    )
+
+
 def run_pipeline(
     df: pd.DataFrame,
-    config: CleanConfig,
+def run_pipeline(  # noqa: PLR0915 - fixed-order pipeline orchestration
     *,
     memory: object | None = None,
     profile: object | None = None,
@@ -69,6 +89,8 @@ def run_pipeline(
     forwarded so the profile backend can replay learned value maps.
     """
     df = _validate_input(df, config)
+    progress_callback = config.progress_callback
+    _emit_progress(progress_callback, "input", "after", df)
     report = CleanReport(
         rows_before=len(df),
         cols_before=df.shape[1],
@@ -84,10 +106,12 @@ def run_pipeline(
         from .context import apply_policy_to_config  # noqa: PLC0415
 
         config = apply_policy_to_config(config, df=df, report=report)
+        _emit_progress(progress_callback, "context", "after", df)
 
     out = df.copy(deep=False) if config.preserve_original else df
     if config.column_names:
         out = normalize_column_names(out, report)
+        _emit_progress(progress_callback, "column_names", "after", out)
 
     # Hard protected-column guard (context policy / mutable=False): fold the
     # protected set into preserve_columns so drop/impute logic honors it, and
@@ -110,44 +134,60 @@ def run_pipeline(
     else:
         guard_snapshot = {}
     out = clean_strings(out, config, report)
+    _emit_progress(progress_callback, "strings", "after", out)
     if config.drop_empty_columns:
         out = drop_empty_columns(out, report, config)
+        _emit_progress(progress_callback, "empty_columns", "after", out)
     if config.drop_empty_rows:
         out = drop_empty_rows(out, report)
+        _emit_progress(progress_callback, "empty_rows", "after", out)
     if config.fix_dtypes:
         out = fix_dtypes(out, config, report)
+        _emit_progress(progress_callback, "dtypes", "after", out)
     if config.drop_constant_columns:
         out = drop_constant_columns(out, config, report)
+        _emit_progress(progress_callback, "constant_columns", "after", out)
     if config.drop_duplicates:
         out = drop_duplicate_rows(out, config, report)
+        _emit_progress(progress_callback, "duplicates", "after", out)
     if config.semantic_enabled:
         # Semantic cleaning runs after representation repair and before the
         # statistical engine, so missing/outlier logic sees repaired values.
         # Lazily imported to keep ``import freshdata`` light.
         from .semantic.apply import run_semantic  # noqa: PLC0415
 
         out = run_semantic(out, config, report, memory=memory, profile=profile)
+        _emit_progress(progress_callback, "semantic", "after", out)
     if config.engine_mode is not None:
         cache = build_engine_cache(out, config)
+        _emit_progress(progress_callback, "engine_cache", "after", out)
         out = auto_missing(
             out, config, report, contexts=cache.contexts, numeric_corr=cache.numeric_corr
         )
+        _emit_progress(progress_callback, "engine_missing", "after", out)
         out = auto_outliers(out, config, report, contexts=cache.contexts)
+        _emit_progress(progress_callback, "engine_outliers", "after", out)
     out = impute_missing(out, config, report)
+    _emit_progress(progress_callback, "missing", "after", out)
     out = handle_outliers(out, config, report)
+    _emit_progress(progress_callback, "outliers", "after", out)
     out = optimize_memory(out, config, report)
+    _emit_progress(progress_callback, "memory", "after", out)
     if guard_snapshot:
         # Physical byte-identity check, before reset_index so row survivors
         # can still be aligned by their original index labels.
         verify_protected(out, guard_snapshot, report)
+        _emit_progress(progress_callback, "protected_columns", "after", out)
     if config.reset_index:
         out = out.reset_index(drop=True)
+        _emit_progress(progress_callback, "index", "after", out)
 
     report.rows_after = len(out)
     report.cols_after = out.shape[1]
     report.memory_after = memory_bytes(out)
     report.missing_after = int(out.isna().sum().sum())
     report.duration_seconds = time.perf_counter() - started
+    _emit_progress(progress_callback, "complete", "after", out)
     return out, report
 
 

src/freshdata/config.py

@@ -11,7 +11,7 @@
 import dataclasses
 import difflib
 import warnings
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 
 _STRATEGY_CHOICES = ("conservative", "balanced", "aggressive", "auto")
@@ -130,6 +130,9 @@ class CleanConfig:
     preserve_original: bool = True
     #: Print a one-line cleaning summary (plus warnings) after each clean.
     verbose: bool = True
+    #: Optional hook called with a small progress event after each enabled
+    #: pipeline stage.
+    progress_callback: Callable[[dict[str, object]], None] | None = None
     #: Columns that must never be dropped by the engine (post-rename names).
     preserve_columns: tuple[str, ...] = ()
     #: The label/target column; never modified by the engine. Columns named
@@ -303,6 +306,8 @@ def __post_init__(self) -> None:
             raise ValueError(f"outlier_factor must be > 0, got {self.outlier_factor!r}")
         if self.sample_size < 1:
             raise ValueError(f"sample_size must be >= 1, got {self.sample_size!r}")
+        if self.progress_callback is not None and not callable(self.progress_callback):
+            raise TypeError("progress_callback must be callable")
         self._validate_semantic()
         self._validate_context()
         extra = _coerce_str_tuple(self.extra_sentinels)

tests/test_api.py

@@ -31,6 +31,45 @@ def test_clean_report_tuple(messy):
     assert len(report) > 0
 
 
+def test_clean_progress_callback_reports_pipeline_stages(messy):
+    events = []
+
+    out = fd.clean(messy, verbose=False, progress_callback=events.append)
+
+    assert isinstance(out, pd.DataFrame)
+    assert events
+    assert all({"step", "status", "rows", "columns"} <= set(event) for event in events)
+
+    steps = [(event["step"], event["status"]) for event in events]
+    assert ("input", "after") in steps
+    assert ("column_names", "after") in steps
+    assert ("strings", "after") in steps
+    assert ("dtypes", "after") in steps
+    assert ("duplicates", "after") in steps
+    assert ("engine_missing", "after") in steps
+    assert ("engine_outliers", "after") in steps
+    assert steps[-1] == ("complete", "after")
+    assert events[-1]["rows"] == len(out)
+    assert events[-1]["columns"] == out.shape[1]
+
+
+def test_cleaner_progress_callback_reports_pipeline_stages(messy):
+    events = []
+    cleaner = fd.Cleaner(verbose=False, progress_callback=events.append)
+
+    out = cleaner.clean(messy)
+
+    assert isinstance(out, pd.DataFrame)
+    assert cleaner.report_ is not None
+    assert [event["step"] for event in events][-1] == "complete"
+
+
+def test_clean_progress_callback_must_be_callable():
+    df = pd.DataFrame({"a": [1]})
+    with pytest.raises(TypeError, match="progress_callback"):
+        fd.clean(df, progress_callback="not-callable")
+
+
 def test_clean_rejects_non_dataframe():
     with pytest.raises(TypeError, match="DataFrame"):
         fd.clean([1, 2, 3])