Skills › Scientific Agent Skills › qutip › _common.py
_common.py
370 lines · skills/qutip/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared bounded-I/O and validation helpers for the QuTiP CLIs."""
4from __future__ import annotations
6import json
7import math
8import os
9import stat
10import sys
11import tempfile
12from collections.abc import Callable, Iterable, Mapping, Sequence
13from pathlib import Path
14from typing import Any
17QUTIP_VERSION = "5.3.1"
18PINNED_INSTALL = 'uv pip install "qutip==5.3.1"'
19DEFAULT_SEED = 20_260_723
21MAX_INPUT_BYTES = 1024 * 1024
22MAX_REPORT_BYTES = 8 * 1024 * 1024
23MAX_HILBERT_DIMENSION = 64
24MAX_SUBSYSTEMS = 8
25MAX_MODEL_OBJECTS = 32
26MAX_TIME_POINTS = 5_001
27MAX_TRAJECTORIES = 2_000
28MAX_SWEEP_RUNS = 6
29MAX_TOTAL_SWEEP_TRAJECTORIES = 4_000
30MAX_FREQUENCY_POINTS = 10_001
31MAX_ABS_FREQUENCY = 100_000.0
32MAX_RATE = 10_000.0
33MAX_TIME = 100_000.0
36class CliError(ValueError):
37 """An expected command-line, schema, or numerical validation error."""
40def load_qutip() -> Any:
41 """Import the exact supported QuTiP release on demand."""
43 try:
44 import qutip
45 except ModuleNotFoundError as exc:
46 if exc.name != "qutip":
47 raise
48 raise CliError(
49 f"QuTiP {QUTIP_VERSION} is required; install it with "
50 f"`{PINNED_INSTALL}`"
51 ) from exc
52 installed = str(getattr(qutip, "__version__", "unknown"))
53 if installed != QUTIP_VERSION:
54 raise CliError(
55 f"this snapshot requires QuTiP {QUTIP_VERSION}, found {installed}; "
56 f"install it with `{PINNED_INSTALL}`"
57 )
58 return qutip
61def _reject_constant(value: str) -> None:
62 raise CliError(f"non-standard JSON constant is not allowed: {value}")
65def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
66 result: dict[str, Any] = {}
67 for key, value in pairs:
68 if key in result:
69 raise CliError(f"duplicate JSON key is not allowed: {key!r}")
70 result[key] = value
71 return result
74def checked_input_file(
75 value: str | os.PathLike[str],
76 *,
77 suffixes: Iterable[str] = (".json",),
78 max_bytes: int = MAX_INPUT_BYTES,
79) -> Path:
80 """Return a bounded regular local input, rejecting URLs and symlinks."""
82 raw = os.fspath(value)
83 if "://" in raw:
84 raise CliError("network URLs are not accepted; provide a local file")
85 path = Path(raw).expanduser()
86 if path.is_symlink():
87 raise CliError(f"input must not be a symlink: {path}")
88 try:
89 info = path.stat()
90 except OSError as exc:
91 raise CliError(f"cannot access input file {path}: {exc}") from exc
92 if not stat.S_ISREG(info.st_mode):
93 raise CliError(f"input is not a regular file: {path}")
94 if info.st_size > max_bytes:
95 raise CliError(f"input is {info.st_size} bytes; limit is {max_bytes}")
96 allowed = {suffix.lower() for suffix in suffixes}
97 if path.suffix.lower() not in allowed:
98 raise CliError(f"input suffix must be one of: {', '.join(sorted(allowed))}")
99 return path.resolve()
102def checked_output_file(
103 value: str | os.PathLike[str],
104 *,
105 force: bool = False,
106) -> Path:
107 """Validate an explicit local JSON destination without following symlinks."""
109 raw = os.fspath(value)
110 if "://" in raw:
111 raise CliError("network URLs are not accepted as output paths")
112 path = Path(raw).expanduser()
113 if path.suffix.lower() != ".json":
114 raise CliError("output path must end in .json")
115 if path.is_symlink():
116 raise CliError(f"output must not be a symlink: {path}")
117 parent = path.parent
118 if not parent.exists() or not parent.is_dir() or parent.is_symlink():
119 raise CliError(f"output parent must be an existing regular directory: {parent}")
120 if path.exists():
121 if not path.is_file():
122 raise CliError(f"output exists and is not a regular file: {path}")
123 if not force:
124 raise CliError(f"refusing to overwrite existing output: {path}")
125 return parent.resolve() / path.name
128def load_json_object(value: str | os.PathLike[str]) -> dict[str, Any]:
129 """Load a bounded strict-JSON object from a local file."""
131 path = checked_input_file(value)
132 try:
133 with path.open("r", encoding="utf-8") as handle:
134 document = json.load(
135 handle,
136 parse_constant=_reject_constant,
137 object_pairs_hook=_unique_object,
138 )
139 except CliError:
140 raise
141 except (OSError, UnicodeError, json.JSONDecodeError) as exc:
142 raise CliError(f"cannot read valid JSON from {path.name}: {exc}") from exc
143 if not isinstance(document, dict):
144 raise CliError("JSON root must be an object")
145 return document
148def validate_keys(
149 value: Mapping[str, Any],
150 *,
151 allowed: Iterable[str],
152 required: Iterable[str] = (),
153 context: str = "object",
154) -> None:
155 """Reject unknown keys and report missing required keys."""
157 allowed_set = set(allowed)
158 required_set = set(required)
159 unknown = sorted(set(value) - allowed_set)
160 missing = sorted(required_set - set(value))
161 if unknown:
162 raise CliError(f"{context} has unknown keys: {', '.join(unknown)}")
163 if missing:
164 raise CliError(f"{context} is missing keys: {', '.join(missing)}")
167def bounded_int(
168 value: Any,
169 *,
170 name: str,
171 minimum: int,
172 maximum: int,
173) -> int:
174 """Validate a bounded integer, rejecting booleans."""
176 if isinstance(value, bool) or not isinstance(value, int):
177 raise CliError(f"{name} must be an integer")
178 if not minimum <= value <= maximum:
179 raise CliError(f"{name} must be from {minimum} through {maximum}")
180 return value
183def finite_float(
184 value: Any,
185 *,
186 name: str,
187 minimum: float | None = None,
188 maximum: float | None = None,
189 minimum_inclusive: bool = True,
190) -> float:
191 """Validate a finite number with optional bounds."""
193 if isinstance(value, bool) or not isinstance(value, (int, float)):
194 raise CliError(f"{name} must be a number")
195 result = float(value)
196 if not math.isfinite(result):
197 raise CliError(f"{name} must be finite")
198 if minimum is not None:
199 valid = result >= minimum if minimum_inclusive else result > minimum
200 if not valid:
201 qualifier = "at least" if minimum_inclusive else "greater than"
202 raise CliError(f"{name} must be {qualifier} {minimum}")
203 if maximum is not None and result > maximum:
204 raise CliError(f"{name} must be no greater than {maximum}")
205 return result
208def bounded_dimensions(value: Any, *, name: str = "dims") -> list[int]:
209 """Validate subsystem dimensions and their bounded product."""
211 if not isinstance(value, list) or not value:
212 raise CliError(f"{name} must be a non-empty JSON array")
213 if len(value) > MAX_SUBSYSTEMS:
214 raise CliError(f"{name} may contain at most {MAX_SUBSYSTEMS} subsystems")
215 dimensions = [
216 bounded_int(item, name=f"{name}[{index}]", minimum=1, maximum=64)
217 for index, item in enumerate(value)
218 ]
219 product = math.prod(dimensions)
220 if product > MAX_HILBERT_DIMENSION:
221 raise CliError(
222 f"{name} product is {product}; limit is {MAX_HILBERT_DIMENSION}"
223 )
224 return dimensions
227def parse_csv_ints(
228 value: str,
229 *,
230 name: str,
231 minimum: int,
232 maximum: int,
233 max_items: int = MAX_SWEEP_RUNS,
234) -> list[int]:
235 """Parse a short comma-separated list of unique increasing integers."""
237 pieces = [piece.strip() for piece in value.split(",")]
238 if not pieces or any(not piece for piece in pieces):
239 raise CliError(f"{name} must be a comma-separated integer list")
240 if len(pieces) > max_items:
241 raise CliError(f"{name} may contain at most {max_items} values")
242 try:
243 values = [int(piece, 10) for piece in pieces]
244 except ValueError as exc:
245 raise CliError(f"{name} must contain only base-10 integers") from exc
246 checked = [
247 bounded_int(item, name=name, minimum=minimum, maximum=maximum)
248 for item in values
249 ]
250 if checked != sorted(set(checked)):
251 raise CliError(f"{name} values must be unique and strictly increasing")
252 return checked
255def ensure_strictly_increasing(values: Sequence[float], *, name: str) -> None:
256 """Require finite strictly increasing values."""
258 if len(values) < 2:
259 raise CliError(f"{name} must contain at least two values")
260 previous = finite_float(values[0], name=f"{name}[0]")
261 for index, value in enumerate(values[1:], start=1):
262 current = finite_float(value, name=f"{name}[{index}]")
263 if current <= previous:
264 raise CliError(f"{name} must be strictly increasing")
265 previous = current
268def to_jsonable(value: Any) -> Any:
269 """Convert trusted result scalars/arrays into strict-JSON-compatible values."""
271 if value is None or isinstance(value, (bool, str, int)):
272 return value
273 if isinstance(value, float):
274 if not math.isfinite(value):
275 raise CliError("report contains a non-finite float")
276 return value
277 if isinstance(value, complex):
278 if not math.isfinite(value.real) or not math.isfinite(value.imag):
279 raise CliError("report contains a non-finite complex value")
280 return {"real": float(value.real), "imag": float(value.imag)}
281 if isinstance(value, Mapping):
282 return {str(key): to_jsonable(item) for key, item in value.items()}
283 if isinstance(value, (list, tuple)):
284 return [to_jsonable(item) for item in value]
285 if hasattr(value, "tolist"):
286 return to_jsonable(value.tolist())
287 if hasattr(value, "item"):
288 return to_jsonable(value.item())
289 return str(value)
292def strict_json_bytes(document: Any) -> bytes:
293 """Serialize deterministic RFC-compatible JSON with a size cap."""
295 normalized = to_jsonable(document)
296 try:
297 payload = (
298 json.dumps(
299 normalized,
300 indent=2,
301 sort_keys=True,
302 ensure_ascii=False,
303 allow_nan=False,
304 )
305 + "\n"
306 ).encode("utf-8")
307 except (TypeError, ValueError) as exc:
308 raise CliError(f"report is not strict JSON: {exc}") from exc
309 if len(payload) > MAX_REPORT_BYTES:
310 raise CliError(
311 f"report is {len(payload)} bytes; limit is {MAX_REPORT_BYTES}"
312 )
313 return payload
316def emit_json(
317 document: Any,
318 *,
319 output: str | os.PathLike[str] | None = None,
320 force: bool = False,
321) -> None:
322 """Print strict JSON or atomically write a private local JSON file."""
324 payload = strict_json_bytes(document)
325 if output is None:
326 print(payload.decode("utf-8"), end="")
327 return
328 destination = checked_output_file(output, force=force)
329 descriptor, temporary_name = tempfile.mkstemp(
330 prefix=f".{destination.name}.",
331 suffix=".tmp",
332 dir=destination.parent,
333 )
334 temporary = Path(temporary_name)
335 try:
336 with os.fdopen(descriptor, "wb") as handle:
337 handle.write(payload)
338 handle.flush()
339 os.fsync(handle.fileno())
340 os.chmod(temporary, 0o600)
341 if destination.exists() and not force:
342 raise CliError(f"refusing to overwrite existing output: {destination}")
343 os.replace(temporary, destination)
344 finally:
345 temporary.unlink(missing_ok=True)
348def add_output_arguments(parser: Any) -> None:
349 """Add common explicit JSON output controls to an ArgumentParser."""
351 parser.add_argument(
352 "--output",
353 help="write strict JSON to this local .json path (default: stdout)",
354 )
355 parser.add_argument(
356 "--force",
357 action="store_true",
358 help="allow replacing an existing regular output file",
359 )
362def run_cli(action: Callable[[], int | None]) -> int:
363 """Run one CLI action with concise expected-error reporting."""
365 try:
366 status = action()
367 except CliError as exc:
368 print(f"error: {exc}", file=sys.stderr)
369 return 2
370 return 0 if status is None else int(status)