Neurokit2

scripts/_common.py

#!/usr/bin/env python3
"""Shared bounded, local-only helpers for the NeuroKit2 skill CLIs."""

from __future__ import annotations

import csv
import hashlib
import io
import json
import math
import os
import stat
import tempfile
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
from typing import Any

NEUROKIT2_VERSION = "0.2.13"
PINNED_INSTALL = 'uv pip install "neurokit2==0.2.13"'
MAX_CSV_BYTES = 64 * 1024 * 1024
MAX_JSON_BYTES = 4 * 1024 * 1024
MAX_OUTPUT_BYTES = 128 * 1024 * 1024
MAX_ROWS = 500_000
MAX_CHANNELS = 64
MAX_SELECTED_COLUMNS = 16
MAX_CELL_CHARS = 4096
MISSING_TOKENS = {"", "na", "n/a", "nan", "none", "null"}


class CliError(ValueError):
    """An expected command-line validation error."""


def _reject_constant(value: str) -> None:
    raise CliError(f"non-standard JSON constant is not allowed: {value}")


def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
    result: dict[str, Any] = {}
    for key, value in pairs:
        if key in result:
            raise CliError(f"duplicate JSON key is not allowed: {key!r}")
        result[key] = value
    return result


def _reject_url(value: str) -> None:
    lowered = value.strip().lower()
    if "://" in lowered or lowered.startswith(
        ("http:", "https:", "ftp:", "s3:", "gs:", "file:")
    ):
        raise CliError("URLs are not accepted; provide a bounded local path")
    if "\x00" in value:
        raise CliError("paths must not contain a NUL byte")


def _absolute_lexical(path: Path) -> Path:
    return Path(os.path.abspath(os.fspath(path)))


def _reject_symlink_components(path: Path) -> None:
    absolute = _absolute_lexical(path)
    current = Path(absolute.anchor)
    for part in absolute.parts[1:]:
        current /= part
        try:
            if current.is_symlink():
                raise CliError(f"symlink paths are not accepted: {current.name!r}")
        except OSError as exc:
            raise CliError(
                f"cannot inspect path component {current.name!r}: {exc}"
            ) from exc


def checked_root(value: str | os.PathLike[str]) -> Path:
    """Return an existing, non-symlink directory used as an I/O boundary."""

    raw = os.fspath(value)
    _reject_url(raw)
    supplied = _absolute_lexical(Path(raw).expanduser())
    if supplied.is_symlink():
        raise CliError("root directory must not itself be a symlink")
    try:
        root = supplied.resolve(strict=True)
        info = root.stat()
    except OSError as exc:
        raise CliError(f"cannot access root directory: {exc}") from exc
    if not stat.S_ISDIR(info.st_mode):
        raise CliError("root must be an existing directory")
    _reject_symlink_components(root)
    return root


def _within_root(candidate: Path, root: Path) -> None:
    try:
        candidate.relative_to(root)
    except ValueError as exc:
        raise CliError("path escapes the declared root directory") from exc


def _suffix_matches(path: Path, suffixes: Iterable[str]) -> bool:
    lowered = path.name.lower()
    return any(lowered.endswith(suffix.lower()) for suffix in suffixes)


def checked_input_file(
    value: str | os.PathLike[str],
    *,
    root: str | os.PathLike[str] = ".",
    suffixes: Iterable[str],
    max_bytes: int,
) -> Path:
    """Return a bounded regular local file within root, rejecting symlinks."""

    raw = os.fspath(value)
    _reject_url(raw)
    root_path = checked_root(root)
    path = Path(raw).expanduser()
    if not path.is_absolute():
        path = root_path / path
    path = _absolute_lexical(path)
    if path.is_symlink():
        raise CliError(f"input must not be a symlink: {path.name!r}")
    try:
        resolved = path.resolve(strict=True)
        info = resolved.stat()
    except OSError as exc:
        raise CliError(f"cannot access input file {path.name!r}: {exc}") from exc
    _within_root(resolved, root_path)
    _reject_symlink_components(resolved)
    if not stat.S_ISREG(info.st_mode):
        raise CliError(f"input is not a regular file: {path.name!r}")
    if info.st_size > max_bytes:
        raise CliError(
            f"input {path.name!r} is {info.st_size} bytes; limit is {max_bytes}"
        )
    if not _suffix_matches(resolved, suffixes):
        allowed = ", ".join(sorted({suffix.lower() for suffix in suffixes}))
        raise CliError(f"input suffix must be one of: {allowed}")
    return resolved


def checked_output_file(
    value: str | os.PathLike[str],
    *,
    root: str | os.PathLike[str] = ".",
    suffixes: Iterable[str],
    force: bool = False,
) -> Path:
    """Return a local output path within root without following symlinks."""

    raw = os.fspath(value)
    _reject_url(raw)
    root_path = checked_root(root)
    path = Path(raw).expanduser()
    if not path.is_absolute():
        path = root_path / path
    path = _absolute_lexical(path)
    if path.name in {"", ".", ".."}:
        raise CliError("output must name a file")
    if not _suffix_matches(path, suffixes):
        allowed = ", ".join(sorted({suffix.lower() for suffix in suffixes}))
        raise CliError(f"output suffix must be one of: {allowed}")
    if path.is_symlink() or path.parent.is_symlink():
        raise CliError("output and its parent must not be symlinks")
    try:
        parent = path.parent.resolve(strict=True)
        parent_info = parent.stat()
    except OSError as exc:
        raise CliError(f"cannot access output parent: {exc}") from exc
    _within_root(parent, root_path)
    _reject_symlink_components(parent)
    if not stat.S_ISDIR(parent_info.st_mode):
        raise CliError("output parent must be an existing directory")
    destination = parent / path.name
    if destination.exists():
        if not destination.is_file():
            raise CliError("output exists and is not a regular file")
        if not force:
            raise CliError(f"refusing to overwrite existing output: {path.name!r}")
    return destination


def atomic_write_bytes(
    destination: str | os.PathLike[str],
    payload: bytes,
    *,
    root: str | os.PathLike[str] = ".",
    suffixes: Iterable[str],
    force: bool = False,
) -> Path:
    """Atomically write a private, bounded file."""

    if len(payload) > MAX_OUTPUT_BYTES:
        raise CliError(f"output is {len(payload)} bytes; limit is {MAX_OUTPUT_BYTES}")
    destination_path = checked_output_file(
        destination, root=root, suffixes=suffixes, force=force
    )
    descriptor, temporary_name = tempfile.mkstemp(
        prefix=f".{destination_path.name}.",
        suffix=".tmp",
        dir=destination_path.parent,
    )
    temporary = Path(temporary_name)
    try:
        with os.fdopen(descriptor, "wb") as handle:
            handle.write(payload)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temporary, 0o600)
        if destination_path.exists() and not force:
            raise CliError(
                f"refusing to overwrite existing output: {destination_path.name!r}"
            )
        os.replace(temporary, destination_path)
    finally:
        temporary.unlink(missing_ok=True)
    return destination_path


def strict_json_bytes(document: Any) -> bytes:
    """Serialize deterministic RFC-compatible JSON."""

    try:
        payload = (
            json.dumps(
                document,
                indent=2,
                sort_keys=True,
                ensure_ascii=False,
                allow_nan=False,
            )
            + "\n"
        ).encode("utf-8")
    except (TypeError, ValueError) as exc:
        raise CliError(f"report is not strict JSON: {exc}") from exc
    if len(payload) > MAX_OUTPUT_BYTES:
        raise CliError(f"report is {len(payload)} bytes; limit is {MAX_OUTPUT_BYTES}")
    return payload


def emit_json(
    document: Any,
    *,
    output: str | os.PathLike[str] | None = None,
    root: str | os.PathLike[str] = ".",
    force: bool = False,
) -> None:
    """Print strict JSON or atomically write it with private permissions."""

    payload = strict_json_bytes(document)
    if output is None:
        print(payload.decode("utf-8"), end="")
        return
    atomic_write_bytes(
        output,
        payload,
        root=root,
        suffixes={".json"},
        force=force,
    )


def load_json_object(
    value: str | os.PathLike[str],
    *,
    root: str | os.PathLike[str] = ".",
    max_bytes: int = MAX_JSON_BYTES,
) -> dict[str, Any]:
    """Load a bounded strict JSON object."""

    path = checked_input_file(
        value,
        root=root,
        suffixes={".json"},
        max_bytes=max_bytes,
    )
    try:
        with path.open("r", encoding="utf-8") as handle:
            document = json.load(
                handle,
                parse_constant=_reject_constant,
                object_pairs_hook=_unique_object,
            )
    except CliError:
        raise
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise CliError(f"cannot read valid JSON from {path.name!r}: {exc}") from exc
    if not isinstance(document, dict):
        raise CliError("JSON root must be an object")
    return document


def validate_keys(
    value: Mapping[str, Any],
    *,
    allowed: Iterable[str],
    required: Iterable[str] = (),
    context: str,
) -> None:
    allowed_set = set(allowed)
    required_set = set(required)
    unknown = sorted(set(value) - allowed_set)
    missing = sorted(required_set - set(value))
    if unknown:
        raise CliError(f"{context} has unknown keys: {', '.join(unknown)}")
    if missing:
        raise CliError(f"{context} is missing keys: {', '.join(missing)}")


def bounded_int(
    value: Any,
    *,
    name: str,
    minimum: int,
    maximum: int,
) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise CliError(f"{name} must be an integer")
    if not minimum <= value <= maximum:
        raise CliError(f"{name} must be between {minimum} and {maximum}")
    return value


def finite_float(
    value: Any,
    *,
    name: str,
    minimum: float | None = None,
    maximum: float | None = None,
) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise CliError(f"{name} must be numeric")
    number = float(value)
    if not math.isfinite(number):
        raise CliError(f"{name} must be finite")
    if minimum is not None and number < minimum:
        raise CliError(f"{name} must be at least {minimum}")
    if maximum is not None and number > maximum:
        raise CliError(f"{name} must be at most {maximum}")
    return number


def parse_name_list(value: str | None, *, name: str) -> list[str]:
    if value is None:
        return []
    names = [item.strip() for item in value.split(",")]
    if not names or any(not item for item in names):
        raise CliError(f"{name} must be a comma-separated list of nonempty names")
    if len(names) != len(set(names)):
        raise CliError(f"{name} must not contain duplicates")
    return names


def require_deidentified(confirmed: bool) -> None:
    if not confirmed:
        raise CliError(
            "local participant data require --deidentified; remove direct identifiers "
            "and review quasi-identifiers before use"
        )


def _read_header(reader: Any) -> list[str]:
    try:
        header = next(reader)
    except StopIteration as exc:
        raise CliError("CSV is empty") from exc
    if not header or any(not name.strip() for name in header):
        raise CliError("CSV header names must be nonempty")
    header = [name.strip() for name in header]
    if len(header) != len(set(header)):
        raise CliError("CSV header names must be unique")
    if len(header) > MAX_CHANNELS:
        raise CliError(f"CSV has {len(header)} columns; limit is {MAX_CHANNELS}")
    return header


def read_numeric_columns(
    path: Path,
    columns: Sequence[str],
    *,
    max_rows: int = MAX_ROWS,
    allow_missing: bool = False,
) -> tuple[dict[str, list[float | None]], list[str], int]:
    """Read selected numeric CSV columns after bounded structural validation."""

    if not 1 <= max_rows <= MAX_ROWS:
        raise CliError(f"max_rows must be between 1 and {MAX_ROWS}")
    if not columns or len(columns) > MAX_SELECTED_COLUMNS:
        raise CliError(f"select between 1 and {MAX_SELECTED_COLUMNS} numeric columns")
    if len(columns) != len(set(columns)):
        raise CliError("selected columns must be unique")
    selected: dict[str, list[float | None]] = {name: [] for name in columns}
    try:
        with path.open("r", encoding="utf-8-sig", newline="") as handle:
            reader = csv.reader(handle)
            header = _read_header(reader)
            missing = [name for name in columns if name not in header]
            if missing:
                raise CliError(f"CSV is missing columns: {', '.join(missing)}")
            positions = {name: header.index(name) for name in columns}
            row_count = 0
            for row_count, row in enumerate(reader, start=1):
                if row_count > max_rows:
                    raise CliError(f"CSV exceeds the row limit of {max_rows}")
                if len(row) != len(header):
                    raise CliError(
                        f"CSV row {row_count + 1} has {len(row)} cells; "
                        f"expected {len(header)}"
                    )
                for name, position in positions.items():
                    cell = row[position].strip()
                    if len(cell) > MAX_CELL_CHARS:
                        raise CliError(
                            f"CSV cell in column {name!r} exceeds "
                            f"{MAX_CELL_CHARS} characters"
                        )
                    if cell.lower() in MISSING_TOKENS:
                        if allow_missing:
                            selected[name].append(None)
                            continue
                        raise CliError(
                            f"column {name!r} contains a missing value at data row "
                            f"{row_count}; segment or apply a documented bounded "
                            "imputation policy before processing"
                        )
                    try:
                        number = float(cell)
                    except ValueError as exc:
                        raise CliError(
                            f"column {name!r} contains non-numeric data at data row "
                            f"{row_count}"
                        ) from exc
                    if not math.isfinite(number):
                        raise CliError(
                            f"column {name!r} contains a non-finite value at data row "
                            f"{row_count}"
                        )
                    selected[name].append(number)
    except CliError:
        raise
    except (OSError, UnicodeError, csv.Error) as exc:
        raise CliError(f"cannot read bounded CSV: {exc}") from exc
    if row_count == 0:
        raise CliError("CSV has a header but no data rows")
    return selected, header, row_count


def _format_csv_value(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, bool):
        return "1" if value else "0"
    if isinstance(value, int):
        return str(value)
    if isinstance(value, float):
        if not math.isfinite(value):
            return ""
        return format(value, ".12g")
    item = getattr(value, "item", None)
    if callable(item):
        converted = item()
        if converted is not value:
            return _format_csv_value(converted)
    text = str(value)
    if len(text) > MAX_CELL_CHARS:
        raise CliError("output CSV cell is too large")
    return text


def csv_bytes(columns: Sequence[str], rows: Iterable[Sequence[Any]]) -> bytes:
    """Serialize a deterministic UTF-8 CSV with Unix newlines."""

    if not columns or len(columns) > MAX_CHANNELS:
        raise CliError(f"CSV output must have 1 to {MAX_CHANNELS} columns")
    if len(columns) != len(set(columns)):
        raise CliError("CSV output columns must be unique")
    buffer = io.StringIO(newline="")
    writer = csv.writer(buffer, lineterminator="\n")
    writer.writerow(columns)
    count = 0
    for count, row in enumerate(rows, start=1):
        if count > MAX_ROWS:
            raise CliError(f"CSV output exceeds the row limit of {MAX_ROWS}")
        values = list(row)
        if len(values) != len(columns):
            raise CliError("CSV output row length does not match columns")
        writer.writerow([_format_csv_value(value) for value in values])
        if buffer.tell() > MAX_OUTPUT_BYTES:
            raise CliError(f"CSV output exceeds {MAX_OUTPUT_BYTES} bytes")
    if count == 0:
        raise CliError("refusing to write a CSV without data rows")
    payload = buffer.getvalue().encode("utf-8")
    if len(payload) > MAX_OUTPUT_BYTES:
        raise CliError(f"CSV output exceeds {MAX_OUTPUT_BYTES} bytes")
    return payload


def write_csv(
    destination: str | os.PathLike[str],
    columns: Sequence[str],
    rows: Iterable[Sequence[Any]],
    *,
    root: str | os.PathLike[str] = ".",
    force: bool = False,
) -> bytes:
    payload = csv_bytes(columns, rows)
    atomic_write_bytes(
        destination,
        payload,
        root=root,
        suffixes={".csv"},
        force=force,
    )
    return payload


def sha256_bytes(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def json_number(value: Any) -> float | int | None:
    """Return a finite JSON scalar from Python or NumPy numeric input."""

    item = getattr(value, "item", None)
    if callable(item):
        value = item()
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, int):
        return value
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if math.isfinite(number) else None


def dataframe_rows(frame: Any) -> Iterable[Sequence[Any]]:
    """Yield rows without importing pandas at module import time."""

    return frame.itertuples(index=False, name=None)


def run_cli(function: Any) -> int:
    """Run a CLI body with concise expected-error handling."""

    try:
        function()
    except CliError as exc:
        print(f"error: {exc}", file=os.sys.stderr)
        return 2
    except ModuleNotFoundError as exc:
        package = exc.name or "a required package"
        print(
            f"error: missing dependency {package!r}; install the pinned environment "
            f"with: {PINNED_INSTALL}",
            file=os.sys.stderr,
        )
        return 2
    except ImportError as exc:
        print(f"error: optional dependency unavailable: {exc}", file=os.sys.stderr)
        return 2
    except KeyboardInterrupt:
        print("error: interrupted", file=os.sys.stderr)
        return 130
    return 0