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 f4141e1..222ae4f 100644 --- a/src/teich/anonymize.py +++ b/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 diff --git a/src/teich/cli.py b/src/teich/cli.py index c88a54a..3777824 100644 --- a/src/teich/cli.py +++ b/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) diff --git a/src/teich/extract.py b/src/teich/extract.py index aa5daab..84839e2 100644 --- a/src/teich/extract.py +++ b/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( diff --git a/tests/test_extract_anonymize_cli.py b/tests/test_extract_anonymize_cli.py index 4a603ed..31eb69e 100644 --- a/tests/test_extract_anonymize_cli.py +++ b/tests/test_extract_anonymize_cli.py @@ -1242,6 +1242,7 @@ 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", @@ -1249,9 +1250,24 @@ def test_extract_inline_anonymization_scrubs_traces_and_leftover_text(tmp_path: 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() 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" },