Skills › Scientific Agent Skills › pytdc › cache_audit.py
cache_audit.py
146 lines · skills/pytdc/scripts/cache_audit.py at 5f2bc03
1#!/usr/bin/env python3
2"""Create a bounded, read-only manifest of a local PyTDC data directory."""
4from __future__ import annotations
6import argparse
7import heapq
8import os
9import sys
10from collections import Counter
11from pathlib import Path
12from typing import Any
14from _common import CliError, bounded_int, emit_json, safe_directory
17def audit_cache(root: Path, *, max_files: int, largest_limit: int) -> dict[str, Any]:
18 """Inspect regular files without following symbolic links."""
20 extension_counts: Counter[str] = Counter()
21 largest: list[tuple[int, str]] = []
22 errors: list[str] = []
23 file_count = 0
24 directory_count = 0
25 symlink_count = 0
26 total_bytes = 0
27 scan_complete = True
29 def on_error(error: OSError) -> None:
30 nonlocal scan_complete
31 scan_complete = False
32 if len(errors) < 20:
33 errors.append(str(error))
35 for current, directories, files in os.walk(
36 root, topdown=True, followlinks=False, onerror=on_error
37 ):
38 directory_count += 1
39 current_path = Path(current)
40 kept_directories: list[str] = []
41 for name in sorted(directories):
42 candidate = current_path / name
43 if candidate.is_symlink():
44 symlink_count += 1
45 else:
46 kept_directories.append(name)
47 directories[:] = kept_directories
49 for name in sorted(files):
50 candidate = current_path / name
51 if candidate.is_symlink():
52 symlink_count += 1
53 continue
54 if file_count >= max_files:
55 scan_complete = False
56 break
57 try:
58 stat = candidate.stat(follow_symlinks=False)
59 except OSError as exc:
60 on_error(exc)
61 continue
62 if not candidate.is_file():
63 continue
65 relative = candidate.relative_to(root).as_posix()
66 size = stat.st_size
67 file_count += 1
68 total_bytes += size
69 extension_counts[candidate.suffix.lower() or "<none>"] += 1
70 item = (size, relative)
71 if largest_limit == 0:
72 continue
73 if len(largest) < largest_limit:
74 heapq.heappush(largest, item)
75 elif item > largest[0]:
76 heapq.heapreplace(largest, item)
77 if not scan_complete and file_count >= max_files:
78 break
80 return {
81 "cache_directory": str(root),
82 "directory_count": directory_count,
83 "download_performed": False,
84 "errors": errors,
85 "extension_counts": dict(sorted(extension_counts.items())),
86 "file_count": file_count,
87 "largest_files": [
88 {"bytes": size, "path": relative}
89 for size, relative in sorted(largest, reverse=True)
90 ],
91 "max_files": max_files,
92 "scan_complete": scan_complete,
93 "symlink_count_skipped": symlink_count,
94 "total_bytes": total_bytes,
95 }
98def build_parser() -> argparse.ArgumentParser:
99 parser = argparse.ArgumentParser(
100 description=(
101 "Audit an existing PyTDC data/oracle directory without network access. "
102 "Symbolic links are skipped and output is bounded."
103 )
104 )
105 parser.add_argument(
106 "--cache-dir",
107 default="data",
108 help="existing relative directory to inspect (default: data)",
109 )
110 parser.add_argument(
111 "--max-files",
112 type=bounded_int(1, 1_000_000),
113 default=100_000,
114 help="stop after this many files (default: 100000)",
115 )
116 parser.add_argument(
117 "--largest",
118 type=bounded_int(0, 100),
119 default=20,
120 help="include at most this many largest files (default: 20)",
121 )
122 parser.add_argument("--output", help="write JSON to a relative workspace path")
123 parser.add_argument(
124 "--force", action="store_true", help="replace an existing --output file"
125 )
126 return parser
129def main(argv: list[str] | None = None) -> int:
130 args = build_parser().parse_args(argv)
131 try:
132 root = safe_directory(
133 args.cache_dir, label="cache directory", must_exist=True
134 )
135 result = audit_cache(
136 root, max_files=args.max_files, largest_limit=args.largest
137 )
138 emit_json(result, args.output, force=args.force)
139 except (CliError, OSError, ValueError) as exc:
140 print(f"error: {exc}", file=sys.stderr)
141 return 2
142 return 0
145if __name__ == "__main__":
146 raise SystemExit(main())