scripts/validation/read_pkl.py

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