From 49e51ce8a446018024196feafc905b81b8195216 Mon Sep 17 00:00:00 2001 From: CTO Agent Date: Wed, 8 Jul 2026 11:58:52 -0400 Subject: [PATCH] feat: show a tqdm progress bar during anonymization 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: Claude Fable 5 --- pyproject.toml | 1 + src/teich/anonymize.py | 20 +++++++++++++++++--- src/teich/cli.py | 27 +++++++++++++++++++++++++-- uv.lock | 2 ++ 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 394d41e..1a8ec3d 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/teich/anonymize.py b/src/teich/anonymize.py index 4df3627..2ba2349 100644 --- a/src/teich/anonymize.py +++ b/src/teich/anonymize.py @@ -10,7 +10,10 @@ import shutil import string from tempfile import NamedTemporaryFile -from typing import Any +from typing import Any, Callable + +# Called after each file with (file_report, files_done, files_total). +AnonymizeProgress = Callable[["AnonymizeFileReport", int, int], None] TEXT_EXTENSIONS = { @@ -55,7 +58,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() @@ -71,13 +80,18 @@ 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 - for source_file in sorted(path for path in input_path.rglob("*") if path.is_file()): + source_files = sorted(path for path in input_path.rglob("*") if path.is_file()) + for source_file in source_files: relative_path = source_file.relative_to(input_path) destination = source_file if in_place else output_path / relative_path 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 diff --git a/src/teich/cli.py b/src/teich/cli.py index 3aa0a8a..816d176 100644 --- a/src/teich/cli.py +++ b/src/teich/cli.py @@ -17,6 +17,8 @@ from rich.table import Table from typer.core import TyperCommand, TyperGroup +from tqdm import tqdm + from .anonymize import anonymize_path from .config import Config from .converter import convert_traces_to_training_data @@ -195,6 +197,27 @@ def _upload_ignore_patterns(cfg: Config) -> list[str]: return patterns +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.""" + 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: + 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(1) + + return anonymize_path(input_path, output_path, in_place=in_place, progress=on_file) + + def _has_non_empty_trace_outputs(traces_dir: Path) -> bool: if not traces_dir.exists(): return False @@ -464,7 +487,7 @@ def _run_extract_command( stale_readme_path = output / "README.md" if stale_readme_path.exists() and stale_readme_path.is_file(): stale_readme_path.unlink() - anonymize_report = None if skip_anonymize else anonymize_path(output, output, in_place=True) + anonymize_report = None if skip_anonymize else _anonymize_with_progress(output, output, in_place=True) readme_path = _write_extract_readme(provider, output, model_filter=model_filter) extracted_message = f"Extracted {result.count} {provider} trace{'s' if result.count != 1 else ''}" if model_filter: @@ -556,7 +579,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) diff --git a/uv.lock b/uv.lock index 16e45d5..91807e2 100644 --- a/uv.lock +++ b/uv.lock @@ -1743,6 +1743,7 @@ dependencies = [ { name = "pydantic" }, { name = "pyyaml" }, { name = "rich" }, + { name = "tqdm" }, { name = "typer" }, { name = "uvicorn" }, { name = "websockets" }, @@ -1774,6 +1775,7 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0" }, { name = "rich", specifier = ">=13.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.4" }, + { name = "tqdm", specifier = ">=4.66" }, { name = "typer", specifier = ">=0.12" }, { name = "uvicorn", specifier = ">=0.29" }, { name = "uvicorn", marker = "extra == 'studio'", specifier = ">=0.29" },