Skills › Scientific Agent Skills › pufferlib › _common.py
_common.py
199 lines · skills/pufferlib/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared safety and strict-JSON helpers for bundled PufferLib CLIs."""
4from __future__ import annotations
6import json
7import math
8import re
9from pathlib import Path
10from typing import Any
12MAX_JSON_BYTES = 1_048_576
13MAX_STEPS = 1_000_000_000
14MAX_ENVS = 65_536
15MAX_WORKERS = 256
16MAX_EVAL_EPISODES = 10_000
18STABLE_SDIST_SHA256 = (
19 "7df3a3e3f5f894d78d2a1f5374097890aec01473183e748abefe4f3faa10eaa9"
20)
21SOURCE_4_COMMIT = "25647630e1b15330bb3153a5a0d3ff8d234c3acf"
23LOGGER_CREDENTIAL_ENV = {
24 "none": None,
25 "wandb": "WANDB_API_KEY",
26 "neptune": "NEPTUNE_API_TOKEN",
27}
29_SLUG = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$")
30_SHA256 = re.compile(r"^[0-9a-f]{64}$")
31_SECRET_KEY = re.compile(
32 r"(?:^|_)(?:api_?key|token|secret|password|credential|"
33 r"private_?key|authorization)(?:$|_)",
34 re.IGNORECASE,
35)
38class UserInputError(ValueError):
39 """Raised for bounded, user-correctable input errors."""
42def bounded_int(value: str | int, *, name: str, minimum: int, maximum: int) -> int:
43 """Parse an integer while rejecting bools and out-of-range values."""
44 if isinstance(value, bool):
45 raise UserInputError(f"{name} must be an integer, not bool")
46 try:
47 parsed = int(value)
48 except (TypeError, ValueError) as exc:
49 raise UserInputError(f"{name} must be an integer") from exc
50 if not minimum <= parsed <= maximum:
51 raise UserInputError(f"{name} must be between {minimum} and {maximum}")
52 return parsed
55def validate_slug(value: str, *, name: str = "name") -> str:
56 """Accept a local identifier, never a dotted import path."""
57 if not isinstance(value, str) or not _SLUG.fullmatch(value):
58 raise UserInputError(
59 f"{name} must match {_SLUG.pattern!r}; dotted import paths are not accepted"
60 )
61 return value
64def validate_sha256(value: str, *, name: str = "sha256") -> str:
65 """Validate a lowercase SHA-256 digest."""
66 if not isinstance(value, str) or not _SHA256.fullmatch(value):
67 raise UserInputError(f"{name} must be 64 lowercase hexadecimal characters")
68 return value
71def _reject_constant(value: str) -> None:
72 raise UserInputError(f"non-finite JSON number is not allowed: {value}")
75def _reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
76 result: dict[str, Any] = {}
77 for key, value in pairs:
78 if key in result:
79 raise UserInputError(f"duplicate JSON key: {key}")
80 result[key] = value
81 return result
84def _assert_finite_json(value: Any, path: str = "$") -> None:
85 if isinstance(value, float) and not math.isfinite(value):
86 raise UserInputError(f"{path} contains a non-finite number")
87 if isinstance(value, dict):
88 for key, item in value.items():
89 _assert_finite_json(item, f"{path}.{key}")
90 elif isinstance(value, list):
91 for index, item in enumerate(value):
92 _assert_finite_json(item, f"{path}[{index}]")
95def strict_json_loads(text: str) -> Any:
96 """Load JSON with duplicate-key and non-finite-number rejection."""
97 try:
98 value = json.loads(
99 text,
100 object_pairs_hook=_reject_duplicate_keys,
101 parse_constant=_reject_constant,
102 )
103 except json.JSONDecodeError as exc:
104 raise UserInputError(
105 f"invalid JSON at line {exc.lineno}, column {exc.colno}: {exc.msg}"
106 ) from exc
107 _assert_finite_json(value)
108 return value
111def strict_json_dumps(value: Any, *, pretty: bool = True) -> str:
112 """Serialize deterministic JSON and reject NaN or Infinity."""
113 _assert_finite_json(value)
114 return json.dumps(
115 value,
116 allow_nan=False,
117 ensure_ascii=True,
118 indent=2 if pretty else None,
119 separators=None if pretty else (",", ":"),
120 sort_keys=True,
121 )
124def emit_json(value: Any, *, pretty: bool = True) -> None:
125 print(strict_json_dumps(value, pretty=pretty))
128def resolve_local_path(
129 value: str | Path,
130 *,
131 root: str | Path,
132 must_exist: bool = True,
133 reject_symlink: bool = True,
134) -> Path:
135 """Resolve a path beneath an explicit root without directory traversal."""
136 root_path = Path(root).expanduser().resolve(strict=True)
137 raw_path = Path(value).expanduser()
138 candidate = raw_path if raw_path.is_absolute() else root_path / raw_path
139 if reject_symlink and candidate.is_symlink():
140 raise UserInputError(f"symlinks are not accepted: {candidate}")
141 try:
142 resolved = candidate.resolve(strict=must_exist)
143 except OSError as exc:
144 raise UserInputError(f"cannot resolve path: {candidate}") from exc
145 try:
146 resolved.relative_to(root_path)
147 except ValueError as exc:
148 raise UserInputError(f"path escapes root {root_path}: {candidate}") from exc
149 return resolved
152def load_json_object(
153 path: str | Path,
154 *,
155 root: str | Path,
156 max_bytes: int = MAX_JSON_BYTES,
157) -> dict[str, Any]:
158 """Read one explicitly named, bounded UTF-8 JSON object."""
159 resolved = resolve_local_path(path, root=root, must_exist=True)
160 size = resolved.stat().st_size
161 if size > max_bytes:
162 raise UserInputError(f"JSON file exceeds {max_bytes} bytes: {resolved}")
163 try:
164 text = resolved.read_text(encoding="utf-8")
165 except (OSError, UnicodeError) as exc:
166 raise UserInputError(f"cannot read UTF-8 JSON file: {resolved}") from exc
167 value = strict_json_loads(text)
168 if not isinstance(value, dict):
169 raise UserInputError("top-level JSON value must be an object")
170 return value
173def secret_key_paths(value: Any, path: str = "$") -> list[str]:
174 """Return key paths that look like credential-bearing configuration."""
175 matches: list[str] = []
176 if isinstance(value, dict):
177 for key, item in value.items():
178 normalized = re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_")
179 child_path = f"{path}.{key}"
180 if _SECRET_KEY.search(normalized):
181 matches.append(child_path)
182 matches.extend(secret_key_paths(item, child_path))
183 elif isinstance(value, list):
184 for index, item in enumerate(value):
185 matches.extend(secret_key_paths(item, f"{path}[{index}]"))
186 return matches
189def require_keys(
190 mapping: dict[str, Any],
191 *,
192 allowed: set[str],
193 required: set[str],
194 path: str,
195) -> list[str]:
196 """Return schema errors for missing and unknown mapping keys."""
197 errors = [f"{path}.{key} is required" for key in sorted(required - mapping.keys())]
198 errors.extend(f"{path}.{key} is not allowed" for key in sorted(mapping.keys() - allowed))
199 return errors