Skills › Scientific Agent Skills › pytdc › _common.py
_common.py
205 lines · skills/pytdc/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared, standard-library-only helpers for the bundled PyTDC CLIs."""
4from __future__ import annotations
6import importlib
7import json
8import math
9import os
10import tempfile
11from importlib.metadata import PackageNotFoundError, version
12from pathlib import Path
13from typing import Any, Iterable, Sequence
16class CliError(ValueError):
17 """A user-facing command-line validation error."""
20def bounded_int(minimum: int, maximum: int):
21 """Return an argparse converter for a bounded integer."""
23 def convert(value: str) -> int:
24 try:
25 parsed = int(value)
26 except ValueError as exc:
27 raise ValueError(f"expected an integer, got {value!r}") from exc
28 if not minimum <= parsed <= maximum:
29 raise ValueError(
30 f"expected an integer from {minimum} through {maximum}, got {parsed}"
31 )
32 return parsed
34 return convert
37def canonical_name(query: str, choices: Iterable[str], label: str) -> str:
38 """Resolve a name case-insensitively without fuzzy matching."""
40 lookup = {str(choice).casefold(): str(choice) for choice in choices}
41 try:
42 return lookup[query.casefold()]
43 except KeyError as exc:
44 preview = ", ".join(sorted(lookup.values())[:20])
45 suffix = "" if len(lookup) <= 20 else ", ..."
46 raise CliError(f"unknown {label} {query!r}; available: {preview}{suffix}") from exc
49def safe_relative_path(
50 raw_path: str,
51 *,
52 label: str,
53 allow_workspace_root: bool = False,
54) -> Path:
55 """Resolve a relative path while preventing writes outside the workspace."""
57 supplied = Path(raw_path)
58 if supplied.is_absolute() or raw_path.startswith("~"):
59 raise CliError(f"{label} must be a relative path inside the current workspace")
61 workspace = Path.cwd().resolve()
62 resolved = (workspace / supplied).resolve(strict=False)
63 if resolved != workspace and workspace not in resolved.parents:
64 raise CliError(f"{label} escapes the current workspace")
65 if resolved == workspace and not allow_workspace_root:
66 raise CliError(f"{label} must not be the workspace root")
67 return resolved
70def safe_directory(
71 raw_path: str,
72 *,
73 label: str,
74 create: bool = False,
75 must_exist: bool = False,
76) -> Path:
77 """Validate a workspace-contained directory."""
79 path = safe_relative_path(raw_path, label=label)
80 if path.exists() and (path.is_symlink() or not path.is_dir()):
81 raise CliError(f"{label} is not a regular directory: {raw_path}")
82 if must_exist and not path.is_dir():
83 raise CliError(f"{label} does not exist: {raw_path}")
84 if create:
85 path.mkdir(parents=True, exist_ok=True)
86 return path
89def safe_input_file(raw_path: str, *, max_bytes: int, label: str) -> Path:
90 """Validate a bounded, workspace-contained input file."""
92 path = safe_relative_path(raw_path, label=label)
93 if path.is_symlink() or not path.is_file():
94 raise CliError(f"{label} is not a regular file: {raw_path}")
95 size = path.stat().st_size
96 if size > max_bytes:
97 raise CliError(f"{label} is {size} bytes; limit is {max_bytes} bytes")
98 return path
101def validate_fractions(values: Sequence[float]) -> tuple[float, float, float]:
102 """Validate train/validation/test fractions."""
104 if len(values) != 3:
105 raise CliError("--frac requires exactly three values")
106 fractions = tuple(float(value) for value in values)
107 if any(not math.isfinite(value) or value < 0 for value in fractions):
108 raise CliError("split fractions must be finite and non-negative")
109 if not math.isclose(sum(fractions), 1.0, rel_tol=0.0, abs_tol=1e-9):
110 raise CliError("split fractions must sum to 1")
111 if fractions[0] <= 0:
112 raise CliError("the training fraction must be greater than zero")
113 return fractions # type: ignore[return-value]
116def load_pytdc_metadata() -> tuple[Any, str]:
117 """Lazily import bundled package metadata without instantiating a dataset."""
119 try:
120 metadata = importlib.import_module("tdc.metadata")
121 package_version = version("PyTDC")
122 except (ImportError, PackageNotFoundError) as exc:
123 raise CliError(
124 "PyTDC is unavailable; install the pinned snapshot with "
125 "`uv pip install \"setuptools==80.9.0\" \"PyTDC==1.1.15\"`"
126 ) from exc
127 return metadata, package_version
130def truncate_value(value: Any, *, max_string: int = 160) -> Any:
131 """Convert common scientific values to bounded JSON-compatible values."""
133 if isinstance(value, Path):
134 return str(value)
135 if isinstance(value, str):
136 if len(value) <= max_string:
137 return value
138 return value[:max_string] + f"...[{len(value) - max_string} chars omitted]"
139 if isinstance(value, dict):
140 return {
141 str(key): truncate_value(item, max_string=max_string)
142 for key, item in value.items()
143 }
144 if isinstance(value, (list, tuple)):
145 return [truncate_value(item, max_string=max_string) for item in value]
146 if isinstance(value, set):
147 return sorted(truncate_value(item, max_string=max_string) for item in value)
148 if hasattr(value, "item"):
149 try:
150 return truncate_value(value.item(), max_string=max_string)
151 except (TypeError, ValueError):
152 pass
153 if hasattr(value, "tolist"):
154 try:
155 return truncate_value(value.tolist(), max_string=max_string)
156 except (TypeError, ValueError):
157 pass
158 return value
161def read_json_file(path: Path) -> Any:
162 """Read strict JSON, rejecting non-standard constants such as NaN."""
164 def reject_constant(value: str) -> None:
165 raise CliError(f"non-standard JSON constant is not allowed: {value}")
167 try:
168 with path.open("r", encoding="utf-8") as handle:
169 return json.load(handle, parse_constant=reject_constant)
170 except json.JSONDecodeError as exc:
171 raise CliError(f"invalid JSON in {path.name}: {exc}") from exc
174def emit_json(payload: Any, output: str | None, *, force: bool = False) -> None:
175 """Print JSON or atomically write it to a safe relative path."""
177 serializable = truncate_value(payload)
178 text = json.dumps(serializable, indent=2, sort_keys=True, allow_nan=False) + "\n"
179 if output is None:
180 print(text, end="")
181 return
183 destination = safe_relative_path(output, label="output path")
184 if destination.exists() and (destination.is_symlink() or destination.is_dir()):
185 raise CliError(f"output path is not a regular file: {output}")
186 if destination.exists() and not force:
187 raise CliError(f"output already exists: {output}; pass --force to replace it")
189 destination.parent.mkdir(parents=True, exist_ok=True)
190 temporary_name: str | None = None
191 try:
192 with tempfile.NamedTemporaryFile(
193 "w",
194 encoding="utf-8",
195 dir=destination.parent,
196 prefix=f".{destination.name}.",
197 suffix=".tmp",
198 delete=False,
199 ) as handle:
200 handle.write(text)
201 temporary_name = handle.name
202 os.replace(temporary_name, destination)
203 finally:
204 if temporary_name is not None:
205 Path(temporary_name).unlink(missing_ok=True)