Skills › Scientific Agent Skills › neurokit2 › _common.py
_common.py
567 lines · skills/neurokit2/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared bounded, local-only helpers for the NeuroKit2 skill CLIs."""
4from __future__ import annotations
6import csv
7import hashlib
8import io
9import json
10import math
11import os
12import stat
13import tempfile
14from collections.abc import Iterable, Mapping, Sequence
15from pathlib import Path
16from typing import Any
18NEUROKIT2_VERSION = "0.2.13"
19PINNED_INSTALL = 'uv pip install "neurokit2==0.2.13"'
20MAX_CSV_BYTES = 64 * 1024 * 1024
21MAX_JSON_BYTES = 4 * 1024 * 1024
22MAX_OUTPUT_BYTES = 128 * 1024 * 1024
23MAX_ROWS = 500_000
24MAX_CHANNELS = 64
25MAX_SELECTED_COLUMNS = 16
26MAX_CELL_CHARS = 4096
27MISSING_TOKENS = {"", "na", "n/a", "nan", "none", "null"}
30class CliError(ValueError):
31 """An expected command-line validation error."""
34def _reject_constant(value: str) -> None:
35 raise CliError(f"non-standard JSON constant is not allowed: {value}")
38def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
39 result: dict[str, Any] = {}
40 for key, value in pairs:
41 if key in result:
42 raise CliError(f"duplicate JSON key is not allowed: {key!r}")
43 result[key] = value
44 return result
47def _reject_url(value: str) -> None:
48 lowered = value.strip().lower()
49 if "://" in lowered or lowered.startswith(
50 ("http:", "https:", "ftp:", "s3:", "gs:", "file:")
51 ):
52 raise CliError("URLs are not accepted; provide a bounded local path")
53 if "\x00" in value:
54 raise CliError("paths must not contain a NUL byte")
57def _absolute_lexical(path: Path) -> Path:
58 return Path(os.path.abspath(os.fspath(path)))
61def _reject_symlink_components(path: Path) -> None:
62 absolute = _absolute_lexical(path)
63 current = Path(absolute.anchor)
64 for part in absolute.parts[1:]:
65 current /= part
66 try:
67 if current.is_symlink():
68 raise CliError(f"symlink paths are not accepted: {current.name!r}")
69 except OSError as exc:
70 raise CliError(
71 f"cannot inspect path component {current.name!r}: {exc}"
72 ) from exc
75def checked_root(value: str | os.PathLike[str]) -> Path:
76 """Return an existing, non-symlink directory used as an I/O boundary."""
78 raw = os.fspath(value)
79 _reject_url(raw)
80 supplied = _absolute_lexical(Path(raw).expanduser())
81 if supplied.is_symlink():
82 raise CliError("root directory must not itself be a symlink")
83 try:
84 root = supplied.resolve(strict=True)
85 info = root.stat()
86 except OSError as exc:
87 raise CliError(f"cannot access root directory: {exc}") from exc
88 if not stat.S_ISDIR(info.st_mode):
89 raise CliError("root must be an existing directory")
90 _reject_symlink_components(root)
91 return root
94def _within_root(candidate: Path, root: Path) -> None:
95 try:
96 candidate.relative_to(root)
97 except ValueError as exc:
98 raise CliError("path escapes the declared root directory") from exc
101def _suffix_matches(path: Path, suffixes: Iterable[str]) -> bool:
102 lowered = path.name.lower()
103 return any(lowered.endswith(suffix.lower()) for suffix in suffixes)
106def checked_input_file(
107 value: str | os.PathLike[str],
108 *,
109 root: str | os.PathLike[str] = ".",
110 suffixes: Iterable[str],
111 max_bytes: int,
112) -> Path:
113 """Return a bounded regular local file within root, rejecting symlinks."""
115 raw = os.fspath(value)
116 _reject_url(raw)
117 root_path = checked_root(root)
118 path = Path(raw).expanduser()
119 if not path.is_absolute():
120 path = root_path / path
121 path = _absolute_lexical(path)
122 if path.is_symlink():
123 raise CliError(f"input must not be a symlink: {path.name!r}")
124 try:
125 resolved = path.resolve(strict=True)
126 info = resolved.stat()
127 except OSError as exc:
128 raise CliError(f"cannot access input file {path.name!r}: {exc}") from exc
129 _within_root(resolved, root_path)
130 _reject_symlink_components(resolved)
131 if not stat.S_ISREG(info.st_mode):
132 raise CliError(f"input is not a regular file: {path.name!r}")
133 if info.st_size > max_bytes:
134 raise CliError(
135 f"input {path.name!r} is {info.st_size} bytes; limit is {max_bytes}"
136 )
137 if not _suffix_matches(resolved, suffixes):
138 allowed = ", ".join(sorted({suffix.lower() for suffix in suffixes}))
139 raise CliError(f"input suffix must be one of: {allowed}")
140 return resolved
143def checked_output_file(
144 value: str | os.PathLike[str],
145 *,
146 root: str | os.PathLike[str] = ".",
147 suffixes: Iterable[str],
148 force: bool = False,
149) -> Path:
150 """Return a local output path within root without following symlinks."""
152 raw = os.fspath(value)
153 _reject_url(raw)
154 root_path = checked_root(root)
155 path = Path(raw).expanduser()
156 if not path.is_absolute():
157 path = root_path / path
158 path = _absolute_lexical(path)
159 if path.name in {"", ".", ".."}:
160 raise CliError("output must name a file")
161 if not _suffix_matches(path, suffixes):
162 allowed = ", ".join(sorted({suffix.lower() for suffix in suffixes}))
163 raise CliError(f"output suffix must be one of: {allowed}")
164 if path.is_symlink() or path.parent.is_symlink():
165 raise CliError("output and its parent must not be symlinks")
166 try:
167 parent = path.parent.resolve(strict=True)
168 parent_info = parent.stat()
169 except OSError as exc:
170 raise CliError(f"cannot access output parent: {exc}") from exc
171 _within_root(parent, root_path)
172 _reject_symlink_components(parent)
173 if not stat.S_ISDIR(parent_info.st_mode):
174 raise CliError("output parent must be an existing directory")
175 destination = parent / path.name
176 if destination.exists():
177 if not destination.is_file():
178 raise CliError("output exists and is not a regular file")
179 if not force:
180 raise CliError(f"refusing to overwrite existing output: {path.name!r}")
181 return destination
184def atomic_write_bytes(
185 destination: str | os.PathLike[str],
186 payload: bytes,
187 *,
188 root: str | os.PathLike[str] = ".",
189 suffixes: Iterable[str],
190 force: bool = False,
191) -> Path:
192 """Atomically write a private, bounded file."""
194 if len(payload) > MAX_OUTPUT_BYTES:
195 raise CliError(f"output is {len(payload)} bytes; limit is {MAX_OUTPUT_BYTES}")
196 destination_path = checked_output_file(
197 destination, root=root, suffixes=suffixes, force=force
198 )
199 descriptor, temporary_name = tempfile.mkstemp(
200 prefix=f".{destination_path.name}.",
201 suffix=".tmp",
202 dir=destination_path.parent,
203 )
204 temporary = Path(temporary_name)
205 try:
206 with os.fdopen(descriptor, "wb") as handle:
207 handle.write(payload)
208 handle.flush()
209 os.fsync(handle.fileno())
210 os.chmod(temporary, 0o600)
211 if destination_path.exists() and not force:
212 raise CliError(
213 f"refusing to overwrite existing output: {destination_path.name!r}"
214 )
215 os.replace(temporary, destination_path)
216 finally:
217 temporary.unlink(missing_ok=True)
218 return destination_path
221def strict_json_bytes(document: Any) -> bytes:
222 """Serialize deterministic RFC-compatible JSON."""
224 try:
225 payload = (
226 json.dumps(
227 document,
228 indent=2,
229 sort_keys=True,
230 ensure_ascii=False,
231 allow_nan=False,
232 )
233 + "\n"
234 ).encode("utf-8")
235 except (TypeError, ValueError) as exc:
236 raise CliError(f"report is not strict JSON: {exc}") from exc
237 if len(payload) > MAX_OUTPUT_BYTES:
238 raise CliError(f"report is {len(payload)} bytes; limit is {MAX_OUTPUT_BYTES}")
239 return payload
242def emit_json(
243 document: Any,
244 *,
245 output: str | os.PathLike[str] | None = None,
246 root: str | os.PathLike[str] = ".",
247 force: bool = False,
248) -> None:
249 """Print strict JSON or atomically write it with private permissions."""
251 payload = strict_json_bytes(document)
252 if output is None:
253 print(payload.decode("utf-8"), end="")
254 return
255 atomic_write_bytes(
256 output,
257 payload,
258 root=root,
259 suffixes={".json"},
260 force=force,
261 )
264def load_json_object(
265 value: str | os.PathLike[str],
266 *,
267 root: str | os.PathLike[str] = ".",
268 max_bytes: int = MAX_JSON_BYTES,
269) -> dict[str, Any]:
270 """Load a bounded strict JSON object."""
272 path = checked_input_file(
273 value,
274 root=root,
275 suffixes={".json"},
276 max_bytes=max_bytes,
277 )
278 try:
279 with path.open("r", encoding="utf-8") as handle:
280 document = json.load(
281 handle,
282 parse_constant=_reject_constant,
283 object_pairs_hook=_unique_object,
284 )
285 except CliError:
286 raise
287 except (OSError, UnicodeError, json.JSONDecodeError) as exc:
288 raise CliError(f"cannot read valid JSON from {path.name!r}: {exc}") from exc
289 if not isinstance(document, dict):
290 raise CliError("JSON root must be an object")
291 return document
294def validate_keys(
295 value: Mapping[str, Any],
296 *,
297 allowed: Iterable[str],
298 required: Iterable[str] = (),
299 context: str,
300) -> None:
301 allowed_set = set(allowed)
302 required_set = set(required)
303 unknown = sorted(set(value) - allowed_set)
304 missing = sorted(required_set - set(value))
305 if unknown:
306 raise CliError(f"{context} has unknown keys: {', '.join(unknown)}")
307 if missing:
308 raise CliError(f"{context} is missing keys: {', '.join(missing)}")
311def bounded_int(
312 value: Any,
313 *,
314 name: str,
315 minimum: int,
316 maximum: int,
317) -> int:
318 if isinstance(value, bool) or not isinstance(value, int):
319 raise CliError(f"{name} must be an integer")
320 if not minimum <= value <= maximum:
321 raise CliError(f"{name} must be between {minimum} and {maximum}")
322 return value
325def finite_float(
326 value: Any,
327 *,
328 name: str,
329 minimum: float | None = None,
330 maximum: float | None = None,
331) -> float:
332 if isinstance(value, bool) or not isinstance(value, (int, float)):
333 raise CliError(f"{name} must be numeric")
334 number = float(value)
335 if not math.isfinite(number):
336 raise CliError(f"{name} must be finite")
337 if minimum is not None and number < minimum:
338 raise CliError(f"{name} must be at least {minimum}")
339 if maximum is not None and number > maximum:
340 raise CliError(f"{name} must be at most {maximum}")
341 return number
344def parse_name_list(value: str | None, *, name: str) -> list[str]:
345 if value is None:
346 return []
347 names = [item.strip() for item in value.split(",")]
348 if not names or any(not item for item in names):
349 raise CliError(f"{name} must be a comma-separated list of nonempty names")
350 if len(names) != len(set(names)):
351 raise CliError(f"{name} must not contain duplicates")
352 return names
355def require_deidentified(confirmed: bool) -> None:
356 if not confirmed:
357 raise CliError(
358 "local participant data require --deidentified; remove direct identifiers "
359 "and review quasi-identifiers before use"
360 )
363def _read_header(reader: Any) -> list[str]:
364 try:
365 header = next(reader)
366 except StopIteration as exc:
367 raise CliError("CSV is empty") from exc
368 if not header or any(not name.strip() for name in header):
369 raise CliError("CSV header names must be nonempty")
370 header = [name.strip() for name in header]
371 if len(header) != len(set(header)):
372 raise CliError("CSV header names must be unique")
373 if len(header) > MAX_CHANNELS:
374 raise CliError(f"CSV has {len(header)} columns; limit is {MAX_CHANNELS}")
375 return header
378def read_numeric_columns(
379 path: Path,
380 columns: Sequence[str],
381 *,
382 max_rows: int = MAX_ROWS,
383 allow_missing: bool = False,
384) -> tuple[dict[str, list[float | None]], list[str], int]:
385 """Read selected numeric CSV columns after bounded structural validation."""
387 if not 1 <= max_rows <= MAX_ROWS:
388 raise CliError(f"max_rows must be between 1 and {MAX_ROWS}")
389 if not columns or len(columns) > MAX_SELECTED_COLUMNS:
390 raise CliError(f"select between 1 and {MAX_SELECTED_COLUMNS} numeric columns")
391 if len(columns) != len(set(columns)):
392 raise CliError("selected columns must be unique")
393 selected: dict[str, list[float | None]] = {name: [] for name in columns}
394 try:
395 with path.open("r", encoding="utf-8-sig", newline="") as handle:
396 reader = csv.reader(handle)
397 header = _read_header(reader)
398 missing = [name for name in columns if name not in header]
399 if missing:
400 raise CliError(f"CSV is missing columns: {', '.join(missing)}")
401 positions = {name: header.index(name) for name in columns}
402 row_count = 0
403 for row_count, row in enumerate(reader, start=1):
404 if row_count > max_rows:
405 raise CliError(f"CSV exceeds the row limit of {max_rows}")
406 if len(row) != len(header):
407 raise CliError(
408 f"CSV row {row_count + 1} has {len(row)} cells; "
409 f"expected {len(header)}"
410 )
411 for name, position in positions.items():
412 cell = row[position].strip()
413 if len(cell) > MAX_CELL_CHARS:
414 raise CliError(
415 f"CSV cell in column {name!r} exceeds "
416 f"{MAX_CELL_CHARS} characters"
417 )
418 if cell.lower() in MISSING_TOKENS:
419 if allow_missing:
420 selected[name].append(None)
421 continue
422 raise CliError(
423 f"column {name!r} contains a missing value at data row "
424 f"{row_count}; segment or apply a documented bounded "
425 "imputation policy before processing"
426 )
427 try:
428 number = float(cell)
429 except ValueError as exc:
430 raise CliError(
431 f"column {name!r} contains non-numeric data at data row "
432 f"{row_count}"
433 ) from exc
434 if not math.isfinite(number):
435 raise CliError(
436 f"column {name!r} contains a non-finite value at data row "
437 f"{row_count}"
438 )
439 selected[name].append(number)
440 except CliError:
441 raise
442 except (OSError, UnicodeError, csv.Error) as exc:
443 raise CliError(f"cannot read bounded CSV: {exc}") from exc
444 if row_count == 0:
445 raise CliError("CSV has a header but no data rows")
446 return selected, header, row_count
449def _format_csv_value(value: Any) -> str:
450 if value is None:
451 return ""
452 if isinstance(value, bool):
453 return "1" if value else "0"
454 if isinstance(value, int):
455 return str(value)
456 if isinstance(value, float):
457 if not math.isfinite(value):
458 return ""
459 return format(value, ".12g")
460 item = getattr(value, "item", None)
461 if callable(item):
462 converted = item()
463 if converted is not value:
464 return _format_csv_value(converted)
465 text = str(value)
466 if len(text) > MAX_CELL_CHARS:
467 raise CliError("output CSV cell is too large")
468 return text
471def csv_bytes(columns: Sequence[str], rows: Iterable[Sequence[Any]]) -> bytes:
472 """Serialize a deterministic UTF-8 CSV with Unix newlines."""
474 if not columns or len(columns) > MAX_CHANNELS:
475 raise CliError(f"CSV output must have 1 to {MAX_CHANNELS} columns")
476 if len(columns) != len(set(columns)):
477 raise CliError("CSV output columns must be unique")
478 buffer = io.StringIO(newline="")
479 writer = csv.writer(buffer, lineterminator="\n")
480 writer.writerow(columns)
481 count = 0
482 for count, row in enumerate(rows, start=1):
483 if count > MAX_ROWS:
484 raise CliError(f"CSV output exceeds the row limit of {MAX_ROWS}")
485 values = list(row)
486 if len(values) != len(columns):
487 raise CliError("CSV output row length does not match columns")
488 writer.writerow([_format_csv_value(value) for value in values])
489 if buffer.tell() > MAX_OUTPUT_BYTES:
490 raise CliError(f"CSV output exceeds {MAX_OUTPUT_BYTES} bytes")
491 if count == 0:
492 raise CliError("refusing to write a CSV without data rows")
493 payload = buffer.getvalue().encode("utf-8")
494 if len(payload) > MAX_OUTPUT_BYTES:
495 raise CliError(f"CSV output exceeds {MAX_OUTPUT_BYTES} bytes")
496 return payload
499def write_csv(
500 destination: str | os.PathLike[str],
501 columns: Sequence[str],
502 rows: Iterable[Sequence[Any]],
503 *,
504 root: str | os.PathLike[str] = ".",
505 force: bool = False,
506) -> bytes:
507 payload = csv_bytes(columns, rows)
508 atomic_write_bytes(
509 destination,
510 payload,
511 root=root,
512 suffixes={".csv"},
513 force=force,
514 )
515 return payload
518def sha256_bytes(payload: bytes) -> str:
519 return hashlib.sha256(payload).hexdigest()
522def json_number(value: Any) -> float | int | None:
523 """Return a finite JSON scalar from Python or NumPy numeric input."""
525 item = getattr(value, "item", None)
526 if callable(item):
527 value = item()
528 if isinstance(value, bool):
529 return int(value)
530 if isinstance(value, int):
531 return value
532 try:
533 number = float(value)
534 except (TypeError, ValueError):
535 return None
536 return number if math.isfinite(number) else None
539def dataframe_rows(frame: Any) -> Iterable[Sequence[Any]]:
540 """Yield rows without importing pandas at module import time."""
542 return frame.itertuples(index=False, name=None)
545def run_cli(function: Any) -> int:
546 """Run a CLI body with concise expected-error handling."""
548 try:
549 function()
550 except CliError as exc:
551 print(f"error: {exc}", file=os.sys.stderr)
552 return 2
553 except ModuleNotFoundError as exc:
554 package = exc.name or "a required package"
555 print(
556 f"error: missing dependency {package!r}; install the pinned environment "
557 f"with: {PINNED_INSTALL}",
558 file=os.sys.stderr,
559 )
560 return 2
561 except ImportError as exc:
562 print(f"error: optional dependency unavailable: {exc}", file=os.sys.stderr)
563 return 2
564 except KeyboardInterrupt:
565 print("error: interrupted", file=os.sys.stderr)
566 return 130
567 return 0