feat: show a tqdm progress bar during anonymization (#26)

anonymize_path accepts an optional per-file progress callback; the CLI
wires it to a tqdm bar (dynamic_ncols, so it resizes with the terminal)
for both teich extract and teich anonymize, showing files processed and
running counts of scrubbed keys, emails, and usernames plus the current
file name. tqdm was already a transitive dependency via
huggingface_hub; it is now declared explicitly.

Co-authored-by: CTO Agent <cto@paperclip.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Arman Rafiee <aran.rafiee@gmail.com>

Author: 4 people

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()