Skills › Scientific Agent Skills › pymatgen › _common.py
_common.py
365 lines · skills/pymatgen/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared, standard-library-first helpers for the bundled pymatgen CLIs."""
4from __future__ import annotations
6import hashlib
7import json
8import math
9import os
10import warnings
11from importlib.metadata import PackageNotFoundError, version
12from pathlib import Path
13from typing import Any
16PYMATGEN_VERSION = "2026.9.24"
17PYMATGEN_CORE_VERSION = "2026.9.23"
18MP_API_VERSION = "0.46.5"
19DEFAULT_MAX_INPUT_BYTES = 50 * 1024 * 1024
20DEFAULT_MAX_OUTPUT_BYTES = 20 * 1024 * 1024
21DEFAULT_MAX_SITES = 10_000
22ABSOLUTE_MAX_INPUT_BYTES = 512 * 1024 * 1024
23ABSOLUTE_MAX_OUTPUT_BYTES = 100 * 1024 * 1024
24ABSOLUTE_MAX_SITES = 100_000
25ABSOLUTE_MAX_PAIRWISE_SITES = 1_000
28class CliError(ValueError):
29 """A user-facing validation or safety error."""
32def reject_url(value: str, label: str = "path") -> None:
33 """Reject URL-like values where a local path is required."""
34 if "://" in value:
35 raise CliError(f"{label} must be a local path, not a URL")
38def checked_input_file(
39 value: str | Path,
40 *,
41 max_bytes: int = DEFAULT_MAX_INPUT_BYTES,
42) -> Path:
43 """Return a resolved, bounded regular input file."""
44 if max_bytes > ABSOLUTE_MAX_INPUT_BYTES:
45 raise CliError(
46 f"input byte limit may not exceed {ABSOLUTE_MAX_INPUT_BYTES}"
47 )
48 raw = str(value)
49 reject_url(raw, "input")
50 path = Path(raw).expanduser()
51 if path.is_symlink():
52 raise CliError("symbolic-link inputs are not accepted")
53 try:
54 resolved = path.resolve(strict=True)
55 except OSError as exc:
56 raise CliError("input file does not exist or cannot be resolved") from exc
57 if not resolved.is_file():
58 raise CliError("input path is not a regular file")
59 size = resolved.stat().st_size
60 if size > max_bytes:
61 raise CliError(f"input exceeds the {max_bytes}-byte limit")
62 return resolved
65def checked_output_file(
66 value: str | Path,
67 *,
68 input_paths: tuple[Path, ...] = (),
69) -> Path:
70 """Validate a new output path without creating or overwriting it."""
71 raw = str(value)
72 reject_url(raw, "output")
73 path = Path(raw).expanduser()
74 if ".." in path.parts:
75 raise CliError("output path may not contain '..'")
76 if path.exists() or path.is_symlink():
77 raise CliError("output already exists; choose a new path")
78 parent = path.parent.resolve(strict=False)
79 if not parent.exists() or not parent.is_dir() or parent.is_symlink():
80 raise CliError("output parent must be an existing, non-symlink directory")
81 resolved = path.resolve(strict=False)
82 for input_path in input_paths:
83 if resolved == input_path.resolve(strict=True):
84 raise CliError("output may not overwrite an input")
85 return resolved
88def _reject_duplicate_pairs(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
89 result: dict[str, Any] = {}
90 for key, value in pairs:
91 if key in result:
92 raise CliError(f"duplicate JSON key: {key!r}")
93 result[key] = value
94 return result
97def _reject_constant(value: str) -> None:
98 raise CliError(f"non-finite JSON number is not allowed: {value}")
101def load_strict_json(
102 value: str | Path,
103 *,
104 max_bytes: int = DEFAULT_MAX_INPUT_BYTES,
105) -> tuple[Any, Path]:
106 """Load bounded JSON while rejecting duplicate keys and NaN/Infinity."""
107 path = checked_input_file(value, max_bytes=max_bytes)
108 try:
109 payload = json.loads(
110 path.read_text(encoding="utf-8"),
111 object_pairs_hook=_reject_duplicate_pairs,
112 parse_constant=_reject_constant,
113 )
114 except UnicodeDecodeError as exc:
115 raise CliError("JSON input must be UTF-8") from exc
116 except json.JSONDecodeError as exc:
117 raise CliError(f"invalid JSON at line {exc.lineno}, column {exc.colno}") from exc
118 return payload, path
121def json_text(payload: Any, *, pretty: bool = True) -> str:
122 """Serialize strict JSON with stable ordering."""
123 return json.dumps(
124 payload,
125 allow_nan=False,
126 ensure_ascii=False,
127 indent=2 if pretty else None,
128 sort_keys=True,
129 )
132def emit_json(payload: Any) -> None:
133 """Print a strict JSON report."""
134 print(json_text(payload))
137def write_text_new(
138 path: Path,
139 text: str,
140 *,
141 max_bytes: int = DEFAULT_MAX_OUTPUT_BYTES,
142) -> None:
143 """Create a UTF-8 text file exclusively after enforcing a byte bound."""
144 if max_bytes > ABSOLUTE_MAX_OUTPUT_BYTES:
145 raise CliError(
146 f"output byte limit may not exceed {ABSOLUTE_MAX_OUTPUT_BYTES}"
147 )
148 encoded = text.encode("utf-8")
149 if len(encoded) > max_bytes:
150 raise CliError(f"output exceeds the {max_bytes}-byte limit")
151 try:
152 with path.open("x", encoding="utf-8", newline="\n") as handle:
153 handle.write(text)
154 if text and not text.endswith("\n"):
155 handle.write("\n")
156 except FileExistsError as exc:
157 raise CliError("output appeared concurrently; nothing was overwritten") from exc
160def write_json_new(
161 path: Path,
162 payload: Any,
163 *,
164 max_bytes: int = DEFAULT_MAX_OUTPUT_BYTES,
165) -> None:
166 """Create a strict JSON file exclusively."""
167 write_text_new(path, json_text(payload), max_bytes=max_bytes)
170def finite_float(value: Any, label: str) -> float:
171 """Parse a finite float."""
172 if isinstance(value, bool):
173 raise CliError(f"{label} must be a finite number")
174 try:
175 number = float(value)
176 except (TypeError, ValueError) as exc:
177 raise CliError(f"{label} must be a finite number") from exc
178 if not math.isfinite(number):
179 raise CliError(f"{label} must be finite")
180 return number
183def positive_int(value: str) -> int:
184 """Argparse converter for positive integers."""
185 try:
186 number = int(value)
187 except ValueError as exc:
188 raise ValueError("must be an integer") from exc
189 if number < 1:
190 raise ValueError("must be at least 1")
191 return number
194def package_versions(names: tuple[str, ...]) -> dict[str, str | None]:
195 """Return installed distribution versions without importing packages."""
196 result: dict[str, str | None] = {}
197 for name in names:
198 try:
199 result[name] = version(name)
200 except PackageNotFoundError:
201 result[name] = None
202 return result
205def sha256_file(path: Path, *, chunk_bytes: int = 1024 * 1024) -> str:
206 """Hash a local file without loading it all into memory."""
207 digest = hashlib.sha256()
208 with path.open("rb") as handle:
209 while chunk := handle.read(chunk_bytes):
210 digest.update(chunk)
211 return digest.hexdigest()
214def structure_oxidation_summary(structure: Any) -> dict[str, Any]:
215 """Summarize oxidation-state decoration without guessing states."""
216 decorated = 0
217 total = 0
218 values: set[float] = set()
219 for site in structure:
220 for specie in site.species:
221 total += 1
222 if hasattr(specie, "oxi_state"):
223 decorated += 1
224 values.add(float(specie.oxi_state))
225 return {
226 "decorated_species_components": decorated,
227 "species_components": total,
228 "all_decorated": bool(total) and decorated == total,
229 "partially_decorated": 0 < decorated < total,
230 "oxidation_states": sorted(values),
231 "guessed": False,
232 }
235def load_structure_json(path: Path, *, max_bytes: int, max_sites: int) -> Any:
236 """Read plain Structure JSON without allowing nested dynamic MSON objects."""
237 from pymatgen.core import Structure
239 payload, _ = load_strict_json(path, max_bytes=max_bytes)
240 if not isinstance(payload, dict):
241 raise CliError("structure JSON must be an object")
242 if payload.get("@module", "pymatgen.core.structure") != "pymatgen.core.structure":
243 raise CliError("structure JSON has an unsupported @module")
244 if payload.get("@class", "Structure") not in {"Structure", "IStructure"}:
245 raise CliError("structure JSON must describe Structure or IStructure")
246 sites = payload.get("sites")
247 if not isinstance(sites, list) or not sites or len(sites) > max_sites:
248 raise CliError("structure JSON must contain a non-empty, bounded sites array")
249 lattice = payload.get("lattice")
250 if not isinstance(lattice, dict) or "matrix" not in lattice:
251 raise CliError("structure JSON requires a lattice matrix")
253 # PeriodicSite.from_dict invokes MontyDecoder on site properties. Only plain
254 # JSON values are accepted there; metadata at the root is checked above.
255 pending = [(value, 1) for key, value in payload.items() if not key.startswith("@")]
256 visited = 0
257 while pending:
258 value, depth = pending.pop()
259 visited += 1
260 if depth > 30 or visited > 1_000_000:
261 raise CliError("structure JSON exceeds nesting or value-count limits")
262 if isinstance(value, dict):
263 if any(key.startswith("@") for key in value):
264 raise CliError("nested MSON metadata is not accepted in structure JSON")
265 pending.extend((item, depth + 1) for item in value.values())
266 elif isinstance(value, list):
267 pending.extend((item, depth + 1) for item in value)
268 elif isinstance(value, float) and not math.isfinite(value):
269 raise CliError("structure JSON contains a non-finite number")
270 try:
271 return Structure.from_dict(payload)
272 except (KeyError, TypeError, ValueError) as exc:
273 raise CliError(f"invalid Structure JSON: {exc}") from exc
276def load_structure(
277 value: str | Path,
278 *,
279 structure_index: int = 0,
280 max_bytes: int = DEFAULT_MAX_INPUT_BYTES,
281 max_sites: int = DEFAULT_MAX_SITES,
282) -> tuple[Any, Path, dict[str, Any]]:
283 """Load one bounded structure and preserve all parser warnings."""
284 if max_sites > ABSOLUTE_MAX_SITES:
285 raise CliError(f"site limit may not exceed {ABSOLUTE_MAX_SITES}")
286 path = checked_input_file(value, max_bytes=max_bytes)
287 warning_messages: list[str] = []
288 parser_messages: list[str] = []
289 structure_count = 1
290 try:
291 with warnings.catch_warnings(record=True) as caught:
292 warnings.simplefilter("always")
293 lower_name = path.name.casefold()
294 if ".yaml" in lower_name or ".yml" in lower_name:
295 raise CliError("bundled intake does not accept YAML; supply CIF or plain Structure JSON")
296 if ".json" in lower_name or ".mson" in lower_name:
297 if path.suffix.casefold() not in {".json", ".mson"}:
298 raise CliError("decompress Structure JSON to a bounded plain .json file first")
299 if structure_index != 0:
300 raise CliError("JSON intake exposes one structure; use index 0")
301 structure = load_structure_json(path, max_bytes=max_bytes, max_sites=max_sites)
302 elif lower_name.endswith((".cif", ".cif.gz", ".cif.bz2", ".mcif")):
303 from pymatgen.io.cif import CifParser
305 parser = CifParser(path)
306 structures = parser.parse_structures(
307 primitive=False,
308 check_occu=True,
309 on_error="raise",
310 )
311 parser_messages = [str(item) for item in parser.warnings]
312 structure_count = len(structures)
313 if not 0 <= structure_index < structure_count:
314 raise CliError(
315 f"structure index {structure_index} is outside "
316 f"0..{structure_count - 1}"
317 )
318 structure = structures[structure_index]
319 else:
320 from pymatgen.core import Structure
322 if structure_index != 0:
323 raise CliError(
324 "non-CIF readers expose one structure; use index 0"
325 )
326 structure = Structure.from_file(path, primitive=False, sort=False)
327 warning_messages = [
328 f"{item.category.__name__}: {item.message}" for item in caught
329 ]
330 except CliError:
331 raise
332 except (OSError, TypeError, ValueError) as exc:
333 raise CliError(
334 f"pymatgen could not parse {path.name}: {type(exc).__name__}: {exc}"
335 ) from exc
336 if len(structure) > max_sites:
337 raise CliError(
338 f"structure has {len(structure)} sites, above the {max_sites}-site limit"
339 )
340 report = {
341 "input_name": path.name,
342 "input_bytes": path.stat().st_size,
343 "structure_index": structure_index,
344 "structures_in_file": structure_count,
345 "python_warnings": warning_messages,
346 "parser_warnings": parser_messages,
347 "warnings_acknowledged": False,
348 }
349 return structure, path, report
352def safe_error_message(exc: BaseException, *, secret: str | None = None) -> str:
353 """Return a bounded error message with a known secret redacted."""
354 message = f"{type(exc).__name__}: {exc}"
355 if secret:
356 message = message.replace(secret, "[REDACTED]")
357 return message[:1000]
360def atomic_link_from_temp(temp_path: Path, output_path: Path) -> None:
361 """Link a completed temporary artifact into place without overwriting."""
362 try:
363 os.link(temp_path, output_path)
364 except FileExistsError as exc:
365 raise CliError("output appeared concurrently; nothing was overwritten") from exc