Coverage for lunar_toolkit/tools.py: 99%
134 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 02:58 +0000
1"""Find and run ISIS / ASP / GDAL tools; keyword options → flags; dry runs."""
3from __future__ import annotations
5import contextlib
6import contextvars
7import json
8import logging
9import os
10import shlex
11import shutil
12import subprocess
13import time
14from collections.abc import Iterator, Mapping, Sequence
15from pathlib import Path
16from typing import Any
18log = logging.getLogger(__name__)
20TOOL_ROOT_ENV_VARS = ("ASP_ROOT", "ISISROOT")
22_DRY_RUN: contextvars.ContextVar[list[list[str]] | None] = contextvars.ContextVar("dry_run", default=None)
23_COMMAND_LOG: contextvars.ContextVar[list[dict[str, Any]] | None] = contextvars.ContextVar(
24 "command_log", default=None
25)
28class ToolError(RuntimeError):
29 """An external tool is missing or exited with a non-zero status."""
31 def __init__(
32 self,
33 message: str,
34 returncode: int | None = None,
35 stdout: str | None = None,
36 stderr: str | None = None,
37 ) -> None:
38 super().__init__(message)
39 self.returncode = returncode
40 self.stdout = stdout
41 self.stderr = stderr
44def find_tool(name: str) -> Path | None:
45 """Path of executable *name*, or None when it cannot be found."""
46 for var in TOOL_ROOT_ENV_VARS:
47 root = os.environ.get(var)
48 if root:
49 p = Path(root) / "bin" / name
50 if p.is_file() and os.access(p, os.X_OK):
51 return p
52 found = shutil.which(name)
53 return Path(found) if found else None
56def tool_env(tool_path: Path) -> dict[str, str]:
57 """Environment for *tool_path* as a conda activation would set it.
59 Unset variables come from the env's ``conda env config vars`` (stored in
60 ``conda-meta/state``), else ``ISISROOT``, ``ISISDATA``, ``PROJ_DATA`` and
61 ``GDAL_DATA`` default to the tool's environment prefix, so tools also run
62 from an env that was not activated (ASP and GDAL fail or crash without them).
63 """
64 env = dict(os.environ)
65 prefix = Path(tool_path).parent.parent
66 for name, setting in _conda_env_vars(prefix).items():
67 env.setdefault(name, setting)
68 defaults = {
69 "ISISROOT": (prefix, prefix / "IsisPreferences"),
70 "ISISDATA": (prefix / "data", prefix / "IsisPreferences"),
71 "PROJ_DATA": (prefix / "share" / "proj", prefix / "share" / "proj" / "proj.db"),
72 "GDAL_DATA": (prefix / "share" / "gdal", prefix / "share" / "gdal"),
73 }
74 for var, (value, marker) in defaults.items():
75 if not env.get(var) and marker.exists():
76 env[var] = str(value)
77 # ASP's Python wrappers (mapproject, parallel_*) run `env python` and
78 # call sibling tools, so the tool's bin dir must lead PATH.
79 bin_dir = str(Path(tool_path).parent)
80 path = env.get("PATH", "")
81 if path.split(os.pathsep)[0] != bin_dir: 81 ↛ 83line 81 didn't jump to line 83 because the condition on line 81 was always true
82 env["PATH"] = os.pathsep.join([bin_dir, path]) if path else bin_dir
83 return env
86def _conda_env_vars(prefix: Path) -> dict[str, str]:
87 try:
88 state = json.loads((prefix / "conda-meta" / "state").read_text())
89 except (OSError, ValueError):
90 return {}
91 env_vars = state.get("env_vars") if isinstance(state, dict) else None
92 return {str(k): str(v) for k, v in env_vars.items()} if isinstance(env_vars, dict) else {}
95def format_command(argv: Sequence[str]) -> str:
96 """*argv* as one command line, quoted for the platform's shell."""
97 return subprocess.list2cmdline(argv) if os.name == "nt" else shlex.join(argv)
100@contextlib.contextmanager
101def dry_run() -> Iterator[list[list[str]]]:
102 """Record commands instead of running them.
104 Yields the list the commands are appended to (tool name first)::
106 with dry_run() as cmds:
107 pipeline.run(...)
108 for c in cmds:
109 print(format_command(c))
110 """
111 cmds: list[list[str]] = []
112 token = _DRY_RUN.set(cmds)
113 try:
114 yield cmds
115 finally:
116 _DRY_RUN.reset(token)
119@contextlib.contextmanager
120def record_commands() -> Iterator[list[dict[str, Any]]]:
121 """Log the commands [`run`][lunar_toolkit.tools.run] executes: ``argv``, ``cwd``, ``started`` (Unix time),
122 ``seconds``, ``returncode``."""
123 log_: list[dict[str, Any]] = []
124 token = _COMMAND_LOG.set(log_)
125 try:
126 yield log_
127 finally:
128 _COMMAND_LOG.reset(token)
131def is_dry_run() -> bool:
132 return _DRY_RUN.get() is not None
135def reuse_existing(path: Path, force: bool = False) -> bool:
136 """Whether a step can be skipped because its output *path* already exists.
138 Parameters
139 ----------
140 path : Path
141 The step's output.
142 force : bool
143 Never reuse; the step is always redone.
145 Returns
146 -------
147 bool
148 True (and logged) when *path* exists and *force* is False.
149 """
150 if Path(path).exists() and not force:
151 log.info("Reusing %s", path)
152 return True
153 return False
156def run(
157 tool: str,
158 args: Sequence[Any] = (),
159 *,
160 cwd: Path | None = None,
161 capture: bool = False,
162 check: bool = True,
163) -> subprocess.CompletedProcess:
164 """Run *tool* with *args*; raise [`ToolError`][lunar_toolkit.tools.ToolError] on failure.
166 With ``capture=True`` stdout/stderr are returned instead of streamed.
167 """
168 argv = [str(a) for a in args]
169 recorded = _DRY_RUN.get()
170 if recorded is not None:
171 recorded.append([tool, *argv])
172 log.info("[dry-run] %s", format_command([tool, *argv]))
173 return subprocess.CompletedProcess([tool, *argv], 0, "", "")
175 path = find_tool(tool)
176 if path is None:
177 raise ToolError(f"{tool} not found; activate the ISIS / ASP environment or set ASP_ROOT / ISISROOT")
178 log.info("▶ %s", tool)
179 log.debug(" cmd: %s", format_command([str(path), *argv]))
180 start = time.time()
181 result = subprocess.run(
182 [str(path), *argv], cwd=cwd, env=tool_env(path), text=True, capture_output=capture
183 )
184 commands = _COMMAND_LOG.get()
185 if commands is not None:
186 commands.append(
187 {
188 "argv": [str(path), *argv],
189 "cwd": str(cwd) if cwd else None,
190 "started": start,
191 "seconds": round(time.time() - start, 3),
192 "returncode": result.returncode,
193 }
194 )
195 if check and result.returncode != 0:
196 tail = (result.stderr or result.stdout or "").strip().splitlines()[-5:]
197 detail = ("\n" + "\n".join(tail)) if capture and tail else " Check output above."
198 raise ToolError(
199 f"{tool} failed (exit {result.returncode}).{detail}",
200 result.returncode,
201 result.stdout,
202 result.stderr,
203 )
204 log.info("✔ %s complete", tool)
205 return result
208def asp_options(options: Mapping[str, Any]) -> list[str]:
209 """Keyword options → ASP-style flags.
211 ``tr=5`` → ``--tr 5``; ``save_sparingly=True`` → ``--save-sparingly``;
212 ``False`` / ``None`` are dropped; tuples and lists expand to several
213 arguments (``t_projwin=(x0, y0, x1, y1)``) while a string with spaces stays
214 one argument (``remove_outliers_params="75 3 100 100"``). Underscores
215 become dashes except in GDAL-style ``t_*`` options (``--t_srs``); single
216 letters get one dash (``o="run/run"`` → ``-o run/run``).
217 """
218 args: list[str] = []
219 for key, value in options.items():
220 if value is None or value is False:
221 continue
222 name = key if key.startswith("t_") else key.replace("_", "-")
223 flag = f"-{name}" if len(name) == 1 else f"--{name}"
224 args += [flag, *_values(value)]
225 return args
228def gdal_options(options: Mapping[str, Any]) -> list[str]:
229 """Keyword options → GDAL-style flags (``a_nodata=0`` → ``-a_nodata 0``)."""
230 args: list[str] = []
231 for key, value in options.items():
232 if value is None or value is False:
233 continue
234 args += [f"-{key}", *_values(value)]
235 return args
238def isis_params(params: Mapping[str, Any]) -> list[str]:
239 """Keyword parameters → ISIS ``key=value`` arguments.
241 A trailing underscore is dropped (``from_=a.cub`` → ``from=a.cub``),
242 booleans become ``true`` / ``false`` and ``None`` is skipped.
243 """
244 args: list[str] = []
245 for key, value in params.items():
246 if value is None:
247 continue
248 if isinstance(value, bool):
249 value = "true" if value else "false"
250 args.append(f"{key.rstrip('_')}={value}")
251 return args
254def _values(value: Any) -> list[str]:
255 if value is True:
256 return []
257 if isinstance(value, (list, tuple)):
258 return [_fmt(v) for v in value]
259 return [_fmt(value)]
262def _fmt(value: Any) -> str:
263 if isinstance(value, float) and value.is_integer():
264 return str(int(value))
265 return str(value)