File size: 2,989 Bytes
0cf9b4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Inspect pickle files in a directory.
Pickle files can contain arbitrary Python objects, so only open files you trust.
"""
from __future__ import annotations
import argparse
import pickle
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any
def _short_repr(value: Any, limit: int = 200) -> str:
text = repr(value).replace("\n", "\\n")
return text if len(text) <= limit else f"{text[: limit - 3]}..."
def _describe(value: Any, *, max_items: int, indent: int = 0) -> list[str]:
"""Return a readable, bounded description of a pickled object."""
prefix = " " * indent
lines = [f"{prefix}type: {type(value).__module__}.{type(value).__qualname__}"]
if hasattr(value, "shape"):
lines.append(f"{prefix}shape: {value.shape}")
if hasattr(value, "dtype"):
lines.append(f"{prefix}dtype: {value.dtype}")
if isinstance(value, Mapping):
lines.append(f"{prefix}length: {len(value)}")
for index, (key, item) in enumerate(value.items()):
if index >= max_items:
lines.append(f"{prefix}... ({len(value) - max_items} more items)")
break
lines.append(f"{prefix}{_short_repr(key)}:")
lines.extend(_describe(item, max_items=max_items, indent=indent + 2))
elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
lines.append(f"{prefix}length: {len(value)}")
for index, item in enumerate(value[:max_items]):
lines.append(f"{prefix}[{index}]:")
lines.extend(_describe(item, max_items=max_items, indent=indent + 2))
if len(value) > max_items:
lines.append(f"{prefix}... ({len(value) - max_items} more items)")
else:
lines.append(f"{prefix}value: {_short_repr(value)}")
return lines
def inspect_file(path: Path, *, max_items: int) -> None:
"""Load and print one pickle file."""
print(f"\n=== {path} ===")
try:
with path.open("rb") as file:
value = pickle.load(file)
except Exception as error: # A bad file should not stop inspection of other files.
print(f"ERROR: {type(error).__name__}: {error}")
return
print("\n".join(_describe(value, max_items=max_items)))
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("directory", type=Path, help="Directory containing pickle files")
parser.add_argument(
"-r",
"--recursive",
action="store_true",
help="Search for pickle files in subdirectories too",
)
parser.add_argument(
"--max-items",
type=int,
default=10,
help="Maximum mapping/sequence items to show per object (default: 10)",
)
args = parser.parse_args()
if not args.directory.is_dir():
parser.error(f"not a directory: {args.directory}")
if args.max_items < 1:
parser.error("--max-items must be at least 1")
pattern = "**/*.pkl" if args.recursive else "*.pkl"
paths = sorted(args.directory.glob(pattern))
if not paths:
print(f"No .pkl files found in {args.directory}")
return
for path in paths:
inspect_file(path, max_items=args.max_items)
if __name__ == "__main__":
main()
|