pyrox.dev
1import argparse
2from functools import partial
3import json
4from multiprocessing import Pool
5import os
6from pathlib import Path
7import subprocess
8
9
10def dropPrefix(path, nixPrefix):
11 return path[len(nixPrefix + "/") :]
12
13
14def processItem(
15 item, nixPrefix, outDir, compression, compressionCommand, compressionExtension
16):
17 narInfoHash = dropPrefix(item["path"], nixPrefix).split("-")[0]
18
19 narFile = outDir / "nar" / f"{narInfoHash}.nar{compressionExtension}"
20 with open(narFile, "wb") as f:
21 subprocess.run(
22 f"nix-store --dump {item['path']} {compressionCommand}",
23 stdout=f,
24 shell=True,
25 check=True,
26 )
27
28 fileHash = (
29 subprocess.run(
30 ["nix-hash", "--base32", "--type", "sha256", "--flat", narFile],
31 capture_output=True,
32 check=True,
33 )
34 .stdout.decode()
35 .strip()
36 )
37 fileSize = os.path.getsize(narFile)
38
39 finalNarFileName = Path("nar") / f"{fileHash}.nar{compressionExtension}"
40 os.rename(narFile, outDir / finalNarFileName)
41
42 with open(outDir / f"{narInfoHash}.narinfo", "wt") as f:
43 f.write(f"StorePath: {item['path']}\n")
44 f.write(f"URL: {finalNarFileName}\n")
45 f.write(f"Compression: {compression}\n")
46 f.write(f"FileHash: sha256:{fileHash}\n")
47 f.write(f"FileSize: {fileSize}\n")
48 f.write(f"NarHash: {item['narHash']}\n")
49 f.write(f"NarSize: {item['narSize']}\n")
50 f.write(
51 f"References: {' '.join(dropPrefix(ref, nixPrefix) for ref in item['references'])}\n"
52 )
53
54
55def main():
56 parser = argparse.ArgumentParser()
57 parser.add_argument("--compression", choices=["none", "xz", "zstd"])
58 args = parser.parse_args()
59
60 compressionCommand = {
61 "none": "",
62 "xz": "| xz -c",
63 "zstd": "| zstd",
64 }[args.compression]
65
66 compressionExtension = {
67 "none": "",
68 "xz": ".xz",
69 "zstd": ".zst",
70 }[args.compression]
71
72 outDir = Path(os.environ["out"])
73 nixPrefix = os.environ["NIX_STORE"]
74 numWorkers = int(os.environ.get("NIX_BUILD_CORES", "4"))
75
76 with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
77 closures = json.load(f)["closure"]
78
79 os.makedirs(outDir / "nar", exist_ok=True)
80
81 with open(outDir / "nix-cache-info", "w") as f:
82 f.write(f"StoreDir: {nixPrefix}\n")
83
84 with Pool(processes=numWorkers) as pool:
85 worker = partial(
86 processItem,
87 nixPrefix=nixPrefix,
88 outDir=outDir,
89 compression=args.compression,
90 compressionCommand=compressionCommand,
91 compressionExtension=compressionExtension,
92 )
93 pool.map(worker, closures)
94
95
96if __name__ == "__main__":
97 main()