Coverage for lunar_toolkit/provenance.py: 95%
179 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 02:58 +0000
1"""Run provenance for SfS work dirs (``provenance.json``): environment, inputs, steps and checks."""
3from __future__ import annotations
5import contextlib
6import hashlib
7import json
8import os
9import platform
10import sys
11import time
12from collections.abc import Iterable, Iterator, Sequence
13from dataclasses import dataclass
14from datetime import datetime, timezone
15from pathlib import Path
16from typing import TYPE_CHECKING, Any
18from .tools import find_tool, is_dry_run, record_commands, run, tool_env
20if TYPE_CHECKING: # pragma: no cover
21 from .config import RunConfig
23SCHEMA_VERSION = 1
24FILE_NAME = "provenance.json"
25MD5_MAX_BYTES = 2 * 1024**3
27# (key, tool, args): first output line is the version.
28VERSION_COMMANDS = (
29 ("asp", "sfs", ["--version"]),
30 ("ale", "isd_generate", ["--version"]),
31 ("gdal", "gdalinfo", ["--version"]),
32)
33KERNEL_GLOBS = {
34 "lsk": "base/kernels/lsk/naif*.tls",
35 "lro_sclk": "lro/kernels/sclk/lro_clkcor_*.tsc",
36}
39def now() -> str:
40 """UTC time as ISO 8601 (seconds)."""
41 return datetime.now(timezone.utc).isoformat(timespec="seconds")
44def _first_line(text: str) -> str | None:
45 lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
46 return lines[0] if lines else None
49def tool_versions() -> dict[str, str | None]:
50 """ASP, ALE, GDAL and ISIS versions (None when a tool is missing); empty under a dry run."""
51 if is_dry_run():
52 return {}
53 out: dict[str, str | None] = {}
54 for key, tool, args in VERSION_COMMANDS:
55 if find_tool(tool) is None:
56 out[key] = None
57 continue
58 res = run(tool, args, capture=True, check=False)
59 out[key] = _first_line(f"{res.stdout or ''}\n{res.stderr or ''}")
60 out["isis"] = isis_version(isis_paths().get("ISISROOT"))
61 return out
64def isis_paths() -> dict[str, str | None]:
65 """``ISISROOT`` / ``ISISDATA`` as the tools see them ([`tool_env`][lunar_toolkit.tools.tool_env])."""
66 tool = find_tool("spiceinit") or find_tool("sfs")
67 env = tool_env(tool) if tool else dict(os.environ)
68 return {k: env.get(k) for k in ("ISISROOT", "ISISDATA")}
71def isis_version(isisroot: str | None) -> str | None:
72 """First field of ``$ISISROOT/isis_version.txt`` (e.g. ``10.0.0``)."""
73 if not isisroot:
74 return None
75 try:
76 text = (Path(isisroot) / "isis_version.txt").read_text()
77 except OSError:
78 return None
79 first = _first_line(text)
80 return first.split("#")[0].strip() if first else None
83def kernel_files(isisdata: str | None) -> dict[str, list[str]]:
84 """Names of the leap-second and LRO clock kernels under *isisdata*."""
85 if not isisdata:
86 return {}
87 root = Path(isisdata)
88 return {k: sorted(p.name for p in root.glob(pattern)) for k, pattern in KERNEL_GLOBS.items()}
91def environment() -> dict[str, Any]:
92 """Package, Python, platform, tool versions, ISIS paths and kernels."""
93 from .constants import APP_VERSION
95 paths = isis_paths()
96 return {
97 "lunar_toolkit": APP_VERSION,
98 "python": platform.python_version(),
99 "platform": platform.platform(),
100 "executable": sys.executable,
101 "tools": tool_versions(),
102 "isis_paths": paths,
103 "kernels": kernel_files(paths.get("ISISDATA")),
104 }
107def md5sum(path: Path, chunk: int = 1 << 20) -> str:
108 """MD5 of *path*, read in *chunk*-byte pieces."""
109 h = hashlib.md5()
110 with open(path, "rb") as f:
111 while block := f.read(chunk):
112 h.update(block)
113 return h.hexdigest()
116def file_record(
117 path: Path, checksum: bool = True, known: dict[str, Any] | None = None, max_bytes: int = MD5_MAX_BYTES
118) -> dict[str, Any]:
119 """``path``, ``size``, ``mtime`` and ``md5`` (None if skipped) of one input file.
121 The MD5 of *known* (an earlier record of the same file) is reused when size
122 and mtime match, so resumed runs do not hash again.
123 """
124 p = Path(path)
125 rec: dict[str, Any] = {"path": str(p), "size": None, "mtime": None, "md5": None}
126 try:
127 st = p.stat()
128 except OSError:
129 return rec
130 rec.update(size=st.st_size, mtime=round(st.st_mtime, 3))
131 if known and known.get("size") == rec["size"] and known.get("mtime") == rec["mtime"]:
132 rec["md5"] = known.get("md5")
133 elif checksum and p.is_file() and st.st_size <= max_bytes:
134 rec["md5"] = md5sum(p)
135 return rec
138def read(path: Path) -> dict[str, Any]:
139 """A ``provenance.json`` (an empty history when missing)."""
140 p = Path(path)
141 if not p.exists():
142 return {"schema": SCHEMA_VERSION, "runs": []}
143 return json.loads(p.read_text())
146def _output_paths(result: Any) -> list[str]:
147 if isinstance(result, Path):
148 return [str(result)] if result.exists() else []
149 if isinstance(result, dict): 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true
150 return _output_paths(list(result.values()))
151 if isinstance(result, (list, tuple)): 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true
152 return [p for r in result for p in _output_paths(r)]
153 path = getattr(result, "mapprojected", None) or getattr(result, "cube", None)
154 return [str(path)] if isinstance(path, Path) and path.exists() else []
157class Provenance:
158 """Append one run (environment, config, inputs, steps) to ``<work_dir>/provenance.json``.
160 Earlier runs are kept, so a resumed work dir holds its whole history. The
161 file is rewritten (atomically) after each step.
162 """
164 def __init__(self, work_dir: Path, config: RunConfig, inputs: Iterable[Path], checksums: bool = True):
165 from .config import config_hash, to_dict
167 self.path = Path(work_dir) / FILE_NAME
168 self.data = read(self.path)
169 known = {r["path"]: r for r in last_run(self.data).get("inputs", [])}
170 self.run: dict[str, Any] = {
171 "started": now(),
172 "finished": None,
173 "status": "running",
174 **environment(),
175 "config_sha256": config_hash(config),
176 "config": to_dict(config),
177 "inputs": [file_record(p, checksums, known.get(str(p))) for p in inputs],
178 "steps": [],
179 }
180 self.data.setdefault("runs", []).append(self.run)
181 self.save()
183 def save(self) -> None:
184 self.path.parent.mkdir(parents=True, exist_ok=True)
185 tmp = self.path.with_suffix(".json.part")
186 tmp.write_text(json.dumps(self.data, indent=1))
187 os.replace(tmp, self.path)
189 @contextlib.contextmanager
190 def step(self, name: str) -> Iterator[dict[str, Any]]:
191 """Record a step: times, the commands [`run`][lunar_toolkit.tools.run] executes and status.
193 Set ``entry["result"]`` inside the block to list its existing output files.
194 """
195 entry: dict[str, Any] = {"name": name, "started": now(), "finished": None, "status": "running"}
196 self.run["steps"].append(entry)
197 self.save()
198 t0 = time.time()
199 with record_commands() as cmds:
200 try:
201 yield entry
202 except BaseException as exc:
203 entry.update(status="error", error=f"{type(exc).__name__}: {exc}")
204 raise
205 else:
206 entry["status"] = "ok"
207 finally:
208 entry.update(
209 finished=now(),
210 seconds=round(time.time() - t0, 3),
211 commands=cmds,
212 outputs=_output_paths(entry.pop("result", None)),
213 )
214 self.save()
216 def finish(self, status: str = "ok") -> None:
217 self.run.update(finished=now(), status=status)
218 self.save()
221def last_run(data: dict[str, Any]) -> dict[str, Any]:
222 runs = data.get("runs") or []
223 return runs[-1] if runs else {}
226@dataclass(frozen=True)
227class Difference:
228 """One mismatch between a recorded run and now: *kind* (``environment``, ``config``, ``input``)."""
230 kind: str
231 key: str
232 recorded: Any
233 current: Any
235 def __str__(self) -> str:
236 return f"{self.kind:11s} {self.key}: {self.recorded!r} -> {self.current!r}"
239def _flatten(d: Any, prefix: str = "") -> dict[str, Any]:
240 if not isinstance(d, dict):
241 return {prefix: d}
242 out: dict[str, Any] = {}
243 for k, v in d.items():
244 out.update(_flatten(v, f"{prefix}.{k}" if prefix else str(k)))
245 return out
248def compare(
249 recorded: dict[str, Any],
250 current_env: dict[str, Any],
251 current_config: dict[str, Any] | None = None,
252 inputs: Sequence[Path] | None = None,
253 checksums: bool = True,
254) -> list[Difference]:
255 """Differences between a recorded run and the current environment, config and inputs.
257 *inputs* default to the recorded input paths; a file whose size or mtime
258 changed is hashed again and reported only when its MD5 differs (or cannot be
259 compared).
260 """
261 diffs: list[Difference] = []
262 keys = ("lunar_toolkit", "python", "tools", "isis_paths", "kernels")
263 rec_env = _flatten({k: recorded.get(k) for k in keys})
264 cur_env = _flatten({k: current_env.get(k) for k in keys})
265 for key in sorted(set(rec_env) | set(cur_env)):
266 if rec_env.get(key) != cur_env.get(key):
267 diffs.append(Difference("environment", key, rec_env.get(key), cur_env.get(key)))
268 if current_config is not None:
269 rec_cfg = _flatten(recorded.get("config") or {})
270 cur_cfg = _flatten(current_config)
271 for key in sorted(set(rec_cfg) | set(cur_cfg)):
272 if rec_cfg.get(key) != cur_cfg.get(key):
273 diffs.append(Difference("config", key, rec_cfg.get(key), cur_cfg.get(key)))
274 known = {r["path"]: r for r in recorded.get("inputs", [])}
275 paths = [str(p) for p in inputs] if inputs is not None else list(known)
276 for p in paths:
277 old = known.get(p)
278 if old is None: 278 ↛ 279line 278 didn't jump to line 279 because the condition on line 278 was never true
279 diffs.append(Difference("input", p, None, "added"))
280 continue
281 new = file_record(Path(p), checksum=False)
282 if new["size"] is None:
283 diffs.append(Difference("input", p, old.get("md5") or old.get("size"), "missing"))
284 elif (new["size"], new["mtime"]) != (old.get("size"), old.get("mtime")):
285 md5 = md5sum(Path(p)) if checksums and old.get("md5") and new["size"] <= MD5_MAX_BYTES else None
286 if md5 is None or md5 != old.get("md5"): 286 ↛ 276line 286 didn't jump to line 276 because the condition on line 286 was always true
287 diffs.append(Difference("input", p, old.get("md5") or old.get("size"), md5 or new["size"]))
288 for p in sorted(set(known) - set(paths)): 288 ↛ 289line 288 didn't jump to line 289 because the loop on line 288 never started
289 diffs.append(Difference("input", p, "used", "not an input now"))
290 return diffs
293def check(
294 work_dir: Path,
295 config: RunConfig | None = None,
296 inputs: Sequence[Path] | None = None,
297 checksums: bool = True,
298) -> list[Difference]:
299 """Compare the last run in ``<work_dir>/provenance.json`` with the current environment.
301 Raises
302 ------
303 FileNotFoundError
304 When the work dir has no provenance yet.
305 """
306 from .config import to_dict
308 path = Path(work_dir) / FILE_NAME
309 if not path.exists():
310 raise FileNotFoundError(f"No {FILE_NAME} in {work_dir}")
311 recorded = last_run(read(path))
312 return compare(recorded, environment(), to_dict(config) if config else None, inputs, checksums)