feat: tqdm progress bar for the anonymization phase

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "rich>=13.0",
     "datasets>=2.19.0",
     "huggingface_hub>=0.23.0",
+    "tqdm>=4.66",
     "fastapi>=0.110",
     "uvicorn>=0.29",
     "websockets>=12",

src/teich/anonymize.py

@@ -13,7 +13,11 @@
 import string
 import sys
 from tempfile import NamedTemporaryFile
-from typing import Any
+from typing import Any, Callable
+
+# Called after each file with (file_report, files_done, files_total). The total
+# is unknown while extraction is discovering and writing traces inline.
+AnonymizeProgress = Callable[["AnonymizeFileReport", int, int | None], None]
 
 
 TEXT_EXTENSIONS = {
@@ -70,7 +74,13 @@ def totals(self) -> dict[str, int]:
         return totals
 
 
-def anonymize_path(input_path: Path, output_path: Path, *, in_place: bool = False) -> AnonymizeReport:
+def anonymize_path(
+    input_path: Path,
+    output_path: Path,
+    *,
+    in_place: bool = False,
+    progress: AnonymizeProgress | None = None,
+) -> AnonymizeReport:
     """Anonymize trace files under input_path."""
     input_path = input_path.expanduser()
     output_path = output_path.expanduser()
@@ -86,6 +96,8 @@ def anonymize_path(input_path: Path, output_path: Path, *, in_place: bool = Fals
         destination = output_path / input_path.name if output_path.exists() and output_path.is_dir() else output_path
         file_report = anonymize_file(input_path, destination)
         report.files.append(file_report)
+        if progress is not None:
+            progress(file_report, 1, 1)
         return report
 
     source_files = sorted(path for path in input_path.rglob("*") if path.is_file())
@@ -95,10 +107,16 @@ def anonymize_path(input_path: Path, output_path: Path, *, in_place: bool = Fals
         # Each file is anonymized independently (fresh TraceAnonymizer per
         # file), so files can be processed in parallel safely.
         with ProcessPoolExecutor(max_workers=workers) as executor:
-            report.files.extend(executor.map(anonymize_file, source_files, destinations))
+            for file_report in executor.map(anonymize_file, source_files, destinations):
+                report.files.append(file_report)
+                if progress is not None:
+                    progress(file_report, len(report.files), len(source_files))
     else:
         for source_file, destination in zip(source_files, destinations):
-            report.files.append(anonymize_file(source_file, destination))
+            file_report = anonymize_file(source_file, destination)
+            report.files.append(file_report)
+            if progress is not None:
+                progress(file_report, len(report.files), len(source_files))
     return report
 
 

src/teich/cli.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from contextlib import contextmanager, nullcontext
 from datetime import datetime, timezone
 import json
 from pathlib import Path
@@ -17,6 +18,8 @@
 from rich.table import Table
 from typer.core import TyperCommand, TyperGroup
 
+from tqdm import tqdm
+
 from .anonymize import anonymize_path
 from .config import CLAUDE_PROVIDER_ALIASES, Config
 from .converter import convert_traces_to_training_data
@@ -195,6 +198,35 @@ def _upload_ignore_patterns(cfg: Config) -> list[str]:
     return patterns
 
 
+@contextmanager
+def _anonymize_progress_bar():
+    """Yield a callback that updates a live file and replacement-count bar."""
+    totals = {"api_key": 0, "email": 0, "username": 0}
+    with tqdm(desc="Anonymizing", unit="file", dynamic_ncols=True, leave=False) as bar:
+
+        def on_file(file_report, done: int, total: int | None) -> None:
+            if total is not None:
+                bar.total = total
+            for key, count in file_report.replacements.items():
+                totals[key] = totals.get(key, 0) + count
+            bar.set_postfix(
+                keys=totals["api_key"],
+                emails=totals["email"],
+                users=totals["username"],
+                file=file_report.path.name,
+                refresh=False,
+            )
+            bar.update(max(0, done - bar.n))
+
+        yield on_file
+
+
+def _anonymize_with_progress(input_path: Path, output_path: Path, *, in_place: bool):
+    """Run anonymize_path with a live tqdm bar showing files and scrub counts."""
+    with _anonymize_progress_bar() as progress:
+        return anonymize_path(input_path, output_path, in_place=in_place, progress=progress)
+
+
 def _has_non_empty_trace_outputs(traces_dir: Path) -> bool:
     if not traces_dir.exists():
         return False
@@ -446,14 +478,17 @@ def _run_extract_command(
     console.print(Panel.fit("Teich Extract", style="bold blue"))
     if provider == "cursor":
         console.print(f"[yellow]{CURSOR_EXTRACTION_NOTICE}[/yellow]", soft_wrap=True)
-    result = extract_local_sessions(
-        provider,
-        output_dir=output,
-        sources=sessions_dir,
-        model_filter=model_filter,
-        clear_destination=True,
-        anonymize=not skip_anonymize,
-    )
+    progress_context = nullcontext(None) if skip_anonymize else _anonymize_progress_bar()
+    with progress_context as anonymize_progress:
+        result = extract_local_sessions(
+            provider,
+            output_dir=output,
+            sources=sessions_dir,
+            model_filter=model_filter,
+            clear_destination=True,
+            anonymize=not skip_anonymize,
+            anonymize_progress=anonymize_progress,
+        )
     if not result.source_paths:
         console.print(f"[red]No local {provider} session folders found.[/red]")
         console.print("[yellow]Pass one or more explicit folders with --sessions-dir.[/yellow]")
@@ -558,7 +593,7 @@ def anonymize(
 ) -> None:
     """Replace emails, home-directory usernames, and API keys with deterministic dummy values."""
     try:
-        report = anonymize_path(input_path, output, in_place=in_place)
+        report = _anonymize_with_progress(input_path, output, in_place=in_place)
     except (FileNotFoundError, ValueError) as exc:
         console.print(f"[red]{exc}[/red]")
         raise typer.Exit(1)

src/teich/extract.py

@@ -15,7 +15,7 @@
 from tempfile import TemporaryDirectory
 from typing import Any, Literal
 
-from .anonymize import TraceAnonymizer, anonymize_file
+from .anonymize import AnonymizeFileReport, AnonymizeProgress, TraceAnonymizer, anonymize_file
 from .tool_schema import CURSOR_BUILTIN_TOOLS
 
 ExtractProvider = Literal["claude", "codex", "cursor", "hermes", "pi"]
@@ -50,22 +50,32 @@ def count(self) -> int:
 class _InlineAnonymizer:
     """Anonymize traces as they are written, one TraceAnonymizer per file."""
 
-    def __init__(self) -> None:
+    def __init__(self, progress: AnonymizeProgress | None = None) -> None:
         self.totals: dict[str, int] = {"email": 0, "username": 0, "api_key": 0}
+        self.progress = progress
+        self.files_done = 0
 
     def _add(self, counts: dict[str, int]) -> None:
         for key, count in counts.items():
             self.totals[key] = self.totals.get(key, 0) + count
 
-    def copy_file(self, source: Path, destination: Path) -> None:
-        report = anonymize_file(source, destination)
+    def _record(self, report: AnonymizeFileReport) -> None:
         self._add(report.replacements)
+        self.files_done += 1
+        if self.progress is not None:
+            self.progress(report, self.files_done, None)
+
+    def copy_file(self, source: Path, destination: Path) -> None:
+        self._record(anonymize_file(source, destination))
 
-    def anonymize_events(self, events: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    def anonymize_events(self, events: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[str, int]]:
         anonymizer = TraceAnonymizer()
         events = [anonymizer.anonymize_value(event) for event in events]
-        self._add(anonymizer.counts)
-        return events
+        replacements = {key: value for key, value in anonymizer.counts.items() if value}
+        return events, replacements
+
+    def record_events_file(self, path: Path, replacements: dict[str, int]) -> None:
+        self._record(AnonymizeFileReport(path=path, output_path=path, replacements=replacements))
 
     def anonymize_remaining_files(self, root: Path, excluded: Iterable[Path]) -> None:
         """Scrub pre-existing output files without re-reading newly written traces."""
@@ -125,6 +135,7 @@ def extract_local_sessions(
     clear_destination: bool = False,
     progress: ProgressCallback | None = None,
     anonymize: bool = False,
+    anonymize_progress: AnonymizeProgress | None = None,
 ) -> ExtractResult:
     """Extract local sessions for provider into output_dir."""
     source_candidates = list(sources) if sources is not None else default_session_sources(provider, home)
@@ -133,7 +144,7 @@ def extract_local_sessions(
     destination_dir.mkdir(parents=True, exist_ok=True)
     if clear_destination:
         _clear_extract_destination(destination_dir)
-    anonymizer = _InlineAnonymizer() if anonymize else None
+    anonymizer = _InlineAnonymizer(anonymize_progress) if anonymize else None
     if provider == "hermes":
         copied_files = _extract_hermes_state_dbs(
             resolved_sources, destination_dir, model_filter=model_filter, anonymizer=anonymizer
@@ -574,14 +585,17 @@ def _write_jsonl_dict_events(
     events: list[dict[str, Any]],
     anonymizer: _InlineAnonymizer | None = None,
 ) -> None:
+    replacements: dict[str, int] = {}
     if anonymizer is not None:
-        events = anonymizer.anonymize_events(events)
+        events, replacements = anonymizer.anonymize_events(events)
     path.parent.mkdir(parents=True, exist_ok=True)
     with path.open("w", encoding="utf-8") as handle:
         for event in events:
             line = json.dumps(event, ensure_ascii=False, separators=(",", ":"))
             line = line.replace("\u0085", "\\u0085").replace("\u2028", "\\u2028").replace("\u2029", "\\u2029")
             handle.write(line + "\n")
+    if anonymizer is not None:
+        anonymizer.record_events_file(path, replacements)
 
 
 def _write_individual_jsonl_rows(

tests/test_extract_anonymize_cli.py

@@ -1242,16 +1242,32 @@ def test_extract_inline_anonymization_scrubs_traces_and_leftover_text(tmp_path:
     stale_text.write_text(f"alice@example.com {original_key}\n", encoding="utf-8")
     stale_binary = output_dir / "archive.bin"
     stale_binary.write_bytes(b"alice@example.com sk-proj-ABCDEFGHIJKLMNOPQRSTUVWXYZ123456")
+    progress_updates: list[tuple[str, int, int | None, dict[str, int]]] = []
 
     result = extract_local_sessions(
         "codex",
         output_dir=output_dir,
         sources=[sessions_dir],
         clear_destination=True,
         anonymize=True,
+        anonymize_progress=lambda report, done, total: progress_updates.append(
+            (report.path.name, done, total, dict(report.replacements))
+        ),
     )
 
     assert result.anonymize_totals == {"email": 3, "username": 2, "api_key": 3}
+    assert [update[0] for update in progress_updates] == [
+        "session-0.jsonl",
+        "session-1.jsonl",
+        "archive.bin",
+        "notes.txt",
+    ]
+    assert [update[1] for update in progress_updates] == [1, 2, 3, 4]
+    assert all(update[2] is None for update in progress_updates)
+    assert {
+        key: sum(update[3].get(key, 0) for update in progress_updates)
+        for key in ("email", "username", "api_key")
+    } == result.anonymize_totals
     for trace in result.copied_files:
         text = trace.read_text(encoding="utf-8")
         assert "/home/user1/project" in text
@@ -1712,6 +1728,51 @@ def test_anonymize_parallel_worker_count_caps_windows(monkeypatch):
     assert anonymize_module._process_worker_count(100) == 61
 
 
+def test_anonymize_path_progress_preserves_parallel_report_order(tmp_path: Path):
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    for index in reversed(range(9)):
+        (input_dir / f"trace-{index}.jsonl").write_text(
+            json.dumps({"message": f"user{index}@example.com"}) + "\n",
+            encoding="utf-8",
+        )
+    updates = []
+
+    report = anonymize_module.anonymize_path(
+        input_dir,
+        tmp_path / "output",
+        progress=lambda file_report, done, total: updates.append((file_report, done, total)),
+    )
+
+    assert [update[0].path.name for update in updates] == [f"trace-{index}.jsonl" for index in range(9)]
+    assert [update[1] for update in updates] == list(range(1, 10))
+    assert [update[2] for update in updates] == [9] * 9
+    assert sum(update[0].replacements.get("email", 0) for update in updates) == report.totals["email"] == 9
+
+
+def test_anonymize_cli_updates_tqdm_progress(tmp_path: Path):
+    source = tmp_path / "trace.jsonl"
+    source.write_text(json.dumps({"message": "alice@example.com"}) + "\n", encoding="utf-8")
+    bar = MagicMock()
+    bar.n = 0
+
+    with patch("teich.cli.tqdm") as mock_tqdm:
+        mock_tqdm.return_value.__enter__.return_value = bar
+        result = runner.invoke(app, ["anonymize", str(source), "--output", str(tmp_path / "output.jsonl")])
+
+    assert result.exit_code == 0, result.output
+    mock_tqdm.assert_called_once_with(desc="Anonymizing", unit="file", dynamic_ncols=True, leave=False)
+    assert bar.total == 1
+    bar.update.assert_called_once_with(1)
+    bar.set_postfix.assert_called_once_with(
+        keys=0,
+        emails=1,
+        users=0,
+        file="trace.jsonl",
+        refresh=False,
+    )
+
+
 def test_anonymize_jsonl_preserves_valid_json_after_escaped_path_replacements(tmp_path: Path):
     input_dir = tmp_path / "input"
     input_dir.mkdir()