solpbc.org
personal memory agent
1#!/usr/bin/env python3
2# SPDX-License-Identifier: AGPL-3.0-only
3# Copyright (c) 2026 sol pbc
4
5"""Repack an unpacked wheel directory after rewriting RECORD."""
6
7from __future__ import annotations
8
9import argparse
10import base64
11import csv
12import hashlib
13import os
14import tempfile
15import zipfile
16from pathlib import Path
17
18
19def _record_hash(path: Path) -> str:
20 digest = hashlib.sha256(path.read_bytes()).digest()
21 encoded = base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
22 return f"sha256={encoded}"
23
24
25def _relative_files(root: Path) -> list[Path]:
26 return sorted(path for path in root.rglob("*") if path.is_file())
27
28
29def _original_file_attrs(wheel_path: Path) -> dict[str, tuple[int, int]]:
30 with zipfile.ZipFile(wheel_path) as wheel:
31 return {
32 info.filename: (info.external_attr, info.create_system)
33 for info in wheel.infolist()
34 if not info.is_dir()
35 }
36
37
38def _rewrite_record(root: Path) -> None:
39 record_paths = list(root.glob("*.dist-info/RECORD"))
40 if len(record_paths) != 1:
41 raise SystemExit(
42 f"expected exactly one *.dist-info/RECORD, found {len(record_paths)}"
43 )
44 record_path = record_paths[0]
45
46 rows: list[list[str]] = []
47 for path in _relative_files(root):
48 arcname = path.relative_to(root).as_posix()
49 if path == record_path:
50 continue
51 rows.append([arcname, _record_hash(path), str(path.stat().st_size)])
52 rows.append([record_path.relative_to(root).as_posix(), "", ""])
53
54 with record_path.open("w", encoding="utf-8", newline="") as handle:
55 csv.writer(handle).writerows(rows)
56
57
58def repack(unpacked_dir: Path, wheel_path: Path) -> None:
59 unpacked_dir = unpacked_dir.resolve()
60 wheel_path = wheel_path.resolve()
61 if not unpacked_dir.is_dir():
62 raise SystemExit(f"unpacked wheel directory does not exist: {unpacked_dir}")
63 if not wheel_path.name.endswith(".whl"):
64 raise SystemExit(f"wheel path must end in .whl: {wheel_path}")
65
66 original_attrs = _original_file_attrs(wheel_path)
67 _rewrite_record(unpacked_dir)
68 fd, tmp_name = tempfile.mkstemp(
69 prefix=f".{wheel_path.name}.", suffix=".tmp", dir=str(wheel_path.parent)
70 )
71 os.close(fd)
72 tmp_path = Path(tmp_name)
73 try:
74 with zipfile.ZipFile(
75 tmp_path, "w", compression=zipfile.ZIP_DEFLATED
76 ) as archive:
77 for path in _relative_files(unpacked_dir):
78 arcname = path.relative_to(unpacked_dir).as_posix()
79 info = zipfile.ZipInfo(arcname)
80 info.compress_type = zipfile.ZIP_DEFLATED
81 if arcname in original_attrs:
82 info.external_attr, info.create_system = original_attrs[arcname]
83 else:
84 info.create_system = 3
85 info.external_attr = (path.stat().st_mode & 0o777) << 16
86 archive.writestr(info, path.read_bytes())
87 os.replace(tmp_path, wheel_path)
88 finally:
89 tmp_path.unlink(missing_ok=True)
90
91
92def main() -> int:
93 parser = argparse.ArgumentParser()
94 parser.add_argument("unpacked_dir", type=Path)
95 parser.add_argument("wheel_path", type=Path)
96 args = parser.parse_args()
97 repack(args.unpacked_dir, args.wheel_path)
98 return 0
99
100
101if __name__ == "__main__":
102 raise SystemExit(main())