xarray-annotated 0.5.2__py3-none-any.whl

xarray_annotated/schema/_annotations.py

@@ -0,0 +1,254 @@
+"""Typed markers declaring a DataArray's structural properties, and the reader.
+
+The structural counterpart to ``units._annotations``: the properties every
+``DataArray`` possesses regardless of physical units — its dimensions
+(``Dims``), coordinates (``Coords``), and dtype (``Dtype``).  Used inside
+``Annotated`` exactly as ``Unit`` is, and — unlike units — **several at once**::
+
+    Annotated[xr.DataArray, Dims("time", "lat", "lon")]
+    Annotated[xr.DataArray, Dims("time", "x"), Dtype("float64")]
+
+Each marker carries the *strictness* of its own check (dims order, exact dtype)
+and an optional per-marker ``on_mismatch`` severity override, because the three
+properties have different natural defaults and a wrong dtype is often less severe
+than a wrong set of dims.  The markers are immutable (``__slots__``) and hashable
+so an annotation carrying them stays hashable/cacheable; ``eq``/``repr`` cover all
+fields, so ``eval(repr(m)) == m``.
+
+There is deliberately **no bare-string shorthand** (a string in the metadata is a
+description, never a declaration): only these typed markers are read.
+"""
+
+from typing import (
+    Annotated,
+    Any,
+    get_args,
+    get_origin,
+)
+
+from .._annotations import _is_dataarray_type, walk_signature
+from ._config import OnMismatch, _validate_on_mismatch
+
+
+def _check_on_mismatch(value: OnMismatch | None) -> OnMismatch | None:
+    """Validate a marker's optional ``on_mismatch`` override at construction."""
+    return None if value is None else _validate_on_mismatch(value)
+
+
+class Dims:
+    """Declare the expected dimension names of a DataArray.
+
+    ``Dims("time", "lat", "lon")`` declares an array over exactly those dims.
+    By default the *set* of dims must match (extra or missing dims fail) but their
+    order is free — xarray operations are order-independent until you drop to
+    numpy.  Pass ``ordered=True`` to also pin the order (e.g. before ``.values``,
+    ``.stack`` or ``apply_ufunc``).
+    """
+
+    _names: tuple[str, ...]
+    _ordered: bool
+    _on_mismatch: OnMismatch | None
+    __slots__ = ("_names", "_ordered", "_on_mismatch")
+
+    def __init__(
+        self,
+        *names: str,
+        ordered: bool = False,
+        on_mismatch: OnMismatch | None = None,
+    ) -> None:
+        self._names = tuple(names)
+        self._ordered = ordered
+        self._on_mismatch = _check_on_mismatch(on_mismatch)
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        """The declared dimension names, in the order given."""
+        return self._names
+
+    @property
+    def ordered(self) -> bool:
+        """Whether dim order (not just the set) must match."""
+        return self._ordered
+
+    @property
+    def on_mismatch(self) -> OnMismatch | None:
+        """Per-marker severity override, or `None` to use the policy default."""
+        return self._on_mismatch
+
+    def __repr__(self) -> str:
+        parts = [repr(n) for n in self._names]
+        if self._ordered:
+            parts.append("ordered=True")
+        if self._on_mismatch is not None:
+            parts.append(f"on_mismatch={self._on_mismatch!r}")
+        return f"Dims({', '.join(parts)})"
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Dims):
+            return (self._names, self._ordered, self._on_mismatch) == (
+                other._names,
+                other._ordered,
+                other._on_mismatch,
+            )
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        return hash((self._names, self._ordered, self._on_mismatch))
+
+
+class Coords:
+    """Declare coordinate variables a DataArray must carry.
+
+    ``Coords("time", "lat")`` declares that those coordinates are present (as
+    *labels*, not merely dims — a dim can exist without coordinate values).
+    Extra coordinates are allowed; only the declared ones must be present.
+    """
+
+    _names: tuple[str, ...]
+    _on_mismatch: OnMismatch | None
+    __slots__ = ("_names", "_on_mismatch")
+
+    def __init__(
+        self,
+        *names: str,
+        on_mismatch: OnMismatch | None = None,
+    ) -> None:
+        self._names = tuple(names)
+        self._on_mismatch = _check_on_mismatch(on_mismatch)
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        """The declared coordinate names."""
+        return self._names
+
+    @property
+    def on_mismatch(self) -> OnMismatch | None:
+        """Per-marker severity override, or `None` to use the policy default."""
+        return self._on_mismatch
+
+    def __repr__(self) -> str:
+        parts = [repr(n) for n in self._names]
+        if self._on_mismatch is not None:
+            parts.append(f"on_mismatch={self._on_mismatch!r}")
+        return f"Coords({', '.join(parts)})"
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Coords):
+            return (self._names, self._on_mismatch) == (
+                other._names,
+                other._on_mismatch,
+            )
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        return hash((self._names, self._on_mismatch))
+
+
+class Dtype:
+    """Declare the expected dtype of a DataArray, e.g. ``Dtype("float64")``.
+
+    By default the check is by *kind*: any float satisfies ``Dtype("float64")``,
+    any integer satisfies ``Dtype("int32")`` — enough to catch an int/float or
+    bool/float mix-up without firing on ``f8`` vs ``f4``.  Pass ``exact=True`` to
+    require the precise dtype (e.g. to pin memory footprint or a typed sink).
+    """
+
+    _dtype: str
+    _exact: bool
+    _on_mismatch: OnMismatch | None
+    __slots__ = ("_dtype", "_exact", "_on_mismatch")
+
+    def __init__(
+        self,
+        dtype: str,
+        *,
+        exact: bool = False,
+        on_mismatch: OnMismatch | None = None,
+    ) -> None:
+        self._dtype = dtype
+        self._exact = exact
+        self._on_mismatch = _check_on_mismatch(on_mismatch)
+
+    @property
+    def dtype(self) -> str:
+        """The declared dtype string."""
+        return self._dtype
+
+    @property
+    def exact(self) -> bool:
+        """Whether the exact dtype (not just its kind) must match."""
+        return self._exact
+
+    @property
+    def on_mismatch(self) -> OnMismatch | None:
+        """Per-marker severity override, or `None` to use the policy default."""
+        return self._on_mismatch
+
+    def __repr__(self) -> str:
+        parts = [repr(self._dtype)]
+        if self._exact:
+            parts.append("exact=True")
+        if self._on_mismatch is not None:
+            parts.append(f"on_mismatch={self._on_mismatch!r}")
+        return f"Dtype({', '.join(parts)})"
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Dtype):
+            return (self._dtype, self._exact, self._on_mismatch) == (
+                other._dtype,
+                other._exact,
+                other._on_mismatch,
+            )
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        return hash((self._dtype, self._exact, self._on_mismatch))
+
+
+#: The marker types the reader recognises in `Annotated` metadata.
+SchemaMarker = Dims | Coords | Dtype
+_SCHEMA_MARKERS = (Dims, Coords, Dtype)
+
+
+def annotated_schema(hint: Any) -> list[SchemaMarker] | None:
+    """Return every schema marker carried by an `Annotated` hint, or `None`.
+
+    Collects *all* `Dims` / `Coords` / `Dtype` markers in the metadata (a hint may
+    declare several structural properties at once), but only when the annotated
+    base type is a `DataArray` (incl. `DataArray | None`).  Non-`Annotated` hints,
+    hints on non-DataArray types, and hints with no schema marker return `None`
+    (so `walk_signature`'s "no declaration" filter drops them).
+
+    Args:
+        hint: A type hint, typically from `get_type_hints(..., include_extras=True)`.
+
+    Returns:
+        A non-empty list of markers in annotation order, or `None`.
+    """
+    if get_origin(hint) is not Annotated:
+        return None
+    args = get_args(hint)
+    if not _is_dataarray_type(args[0]):
+        return None
+    markers = [m for m in args[1:] if isinstance(m, _SCHEMA_MARKERS)]
+    return markers or None
+
+
+def schema_from_signature(
+    func: object,
+) -> tuple[
+    dict[str, list[SchemaMarker]],
+    dict[str, list[SchemaMarker]] | list[SchemaMarker] | None,
+]:
+    """Extract declared schema markers from a function's type annotations.
+
+    Args:
+        func: A callable whose hints carry `Annotated[DataArray, <markers>]`.
+
+    Returns:
+        An `(inputs, output)` pair, where each declaration is the list of markers
+        on that parameter/field; `output` is a per-field dict for a `TypedDict` or
+        dataclass return, a single list for one declared `DataArray` return, or
+        `None`.
+    """
+    return walk_signature(func, annotated_schema)

xarray_annotated/schema/_check.py

@@ -0,0 +1,232 @@
+"""Runtime structural validation of a DataArray against declared markers.
+
+The "usage side" of the schema domain: given a `DataArray` and the `Dims` /
+`Coords` / `Dtype` markers declared for it, check each one under the active policy.
+Unlike the units domain this **never mutates** — it asserts and (per `on_mismatch`)
+raises, warns, or ignores, then returns the array unchanged, so a decorator can
+drop it into an argument-validation pass symmetrically with `check_units`.
+
+The per-property checks are independent functions dispatched by marker type, so a
+new structural property is added by writing one checker and one dispatch entry —
+nothing else here changes.
+"""
+
+import warnings
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import xarray as xr
+
+from ._annotations import Coords, Dims, Dtype, SchemaMarker
+from ._config import OnMismatch, _validate_on_mismatch, get_policy
+
+
+class SchemaWarning(UserWarning):
+    """Warning issued for a structural mismatch under ``on_mismatch="warn"``."""
+
+
+class SchemaError(Exception):
+    """Raised for a structural mismatch under ``on_mismatch="error"``.
+
+    Deliberately **not** a ``ValueError``: a structural mismatch (the *data* fails
+    the declaration) is a distinct event from a malformed *declaration* (an
+    unparseable dtype, duplicate dim names), which raises ``ValueError`` from
+    ``assert_valid_schema``. Keeping them on separate hierarchies means catching
+    one never silently swallows the other.
+    """
+
+
+def assert_valid_schema(marker: SchemaMarker, context: str) -> None:
+    """Validate a marker *declaration* itself, independent of any DataArray.
+
+    Used at decoration time to fail fast on a malformed declaration (e.g. an
+    unparseable dtype string or duplicate dim names) rather than only when the
+    decorated function first runs.
+
+    Args:
+        marker: The declared `Dims`, `Coords`, or `Dtype` marker.
+        context: A label for the error message (e.g. ``"f input 'x'"``).
+
+    Raises:
+        ValueError: If the declaration is malformed.
+    """
+    kind = type(marker).__name__
+    if isinstance(marker, Dtype):
+        try:
+            np.dtype(marker.dtype)
+        except TypeError as exc:
+            raise ValueError(f"{context}: invalid dtype {marker.dtype!r}.") from exc
+    else:  # Dims or Coords: a tuple of names
+        for name in marker.names:
+            if not isinstance(name, str) or not name:
+                raise ValueError(
+                    f"{context}: {kind} names must be non-empty strings, "
+                    f"got {marker.names!r}."
+                )
+        if len(set(marker.names)) != len(marker.names):
+            raise ValueError(f"{context}: {kind} has duplicate names {marker.names!r}.")
+
+
+def _check_dims(da: xr.DataArray, marker: Dims) -> tuple[bool, str]:
+    """Check a DataArray's dims against a `Dims` marker. Returns (ok, detail)."""
+    actual = tuple(da.dims)
+    if marker.ordered:
+        ok = actual == marker.names
+        detail = f"dims order mismatch: expected {marker.names}, got {actual}"
+    else:
+        ok = set(actual) == set(marker.names)
+        detail = f"dims mismatch: expected {marker.names} in any order, got {actual}"
+    return ok, detail
+
+
+def _check_coords(da: xr.DataArray, marker: Coords) -> tuple[bool, str]:
+    """Check a DataArray carries the declared coords. Returns (ok, detail)."""
+    actual = set(map(str, da.coords))
+    missing = [n for n in marker.names if n not in actual]
+    detail = f"missing coords {missing}; present: {sorted(actual)}"
+    return not missing, detail
+
+
+def _check_dtype(da: xr.DataArray, marker: Dtype) -> tuple[bool, str]:
+    """Check a DataArray's dtype against a `Dtype` marker. Returns (ok, detail)."""
+    declared = np.dtype(marker.dtype)
+    actual = da.dtype
+    if marker.exact:
+        ok = actual == declared
+        detail = f"dtype mismatch: expected exactly {declared}, got {actual}"
+    else:
+        ok = actual.kind == declared.kind
+        detail = (
+            f"dtype kind mismatch: expected kind {declared.kind!r} "
+            f"(like {declared}), got {actual}"
+        )
+    return ok, detail
+
+
+def dims_compatible(a: Dims, b: Dims) -> bool:
+    """Return whether two `Dims` declarations can describe the same array.
+
+    The marker-vs-marker counterpart to validating an array against a `Dims`
+    declaration (no array in hand), for a static / build-time check of a
+    producer/consumer edge: two declarations are
+    *provably inconsistent* only if their dim **sets** differ, or if both pin the
+    order (`ordered=True`) and the orders disagree.  A loose declaration on either
+    side can always be satisfied, so it never conflicts.
+
+    Args:
+        a: A `Dims` marker.
+        b: Another `Dims` marker.
+
+    Returns:
+        `True` unless the two are provably inconsistent.
+
+    Examples:
+        >>> from xarray_annotated.schema import Dims, dims_compatible
+        >>> dims_compatible(Dims("x", "y"), Dims("y", "x"))
+        True
+        >>> dims_compatible(Dims("x"), Dims("x", "y"))
+        False
+        >>> dims_compatible(Dims("x", "y", ordered=True), Dims("y", "x", ordered=True))
+        False
+    """
+    if set(a.names) != set(b.names):
+        return False
+    return not (a.ordered and b.ordered and a.names != b.names)
+
+
+def dtype_compatible(a: Dtype, b: Dtype) -> bool:
+    """Return whether two `Dtype` declarations can describe the same array.
+
+    The marker-vs-marker counterpart to validating an array against a `Dtype`
+    declaration (no array in hand): two
+    declarations are *provably inconsistent* only if their numpy **kinds** differ
+    (float vs int), or if both require the exact dtype (`exact=True`) and those
+    dtypes differ (`f8` vs `f4`).  A kind-only declaration matches any width, so
+    it never conflicts on width alone.
+
+    Args:
+        a: A `Dtype` marker.
+        b: Another `Dtype` marker.
+
+    Returns:
+        `True` unless the two are provably inconsistent.
+
+    Examples:
+        >>> from xarray_annotated.schema import Dtype, dtype_compatible
+        >>> dtype_compatible(Dtype("float64"), Dtype("float32"))
+        True
+        >>> dtype_compatible(Dtype("float64"), Dtype("int64"))
+        False
+        >>> dtype_compatible(Dtype("float64", exact=True), Dtype("float32", exact=True))
+        False
+    """
+    if np.dtype(a.dtype).kind != np.dtype(b.dtype).kind:
+        return False
+    return not (a.exact and b.exact and np.dtype(a.dtype) != np.dtype(b.dtype))
+
+
+# There is deliberately no `coords_compatible`: two `Coords` declarations are
+# lower bounds (each merely requires its names be present, extras allowed), so
+# they can never be *proven* inconsistent — a build-time edge check skips coords.
+
+
+_CHECKERS: dict[type, Callable[[xr.DataArray, Any], tuple[bool, str]]] = {
+    Dims: _check_dims,
+    Coords: _check_coords,
+    Dtype: _check_dtype,
+}
+
+
+def _report(ok: bool, severity: OnMismatch, message: str) -> None:
+    """Act on a check result under the resolved severity."""
+    if ok or severity == "ignore":
+        return
+    if severity == "error":
+        raise SchemaError(message)
+    # stacklevel=3 targets the declare_schema call path (user -> wrapper ->
+    # check_schema -> _report); one frame too far for a direct check_schema call.
+    warnings.warn(message, SchemaWarning, stacklevel=3)
+
+
+def check_schema(
+    da: xr.DataArray,
+    declared: SchemaMarker | list[SchemaMarker],
+    name: str,
+    on_mismatch: OnMismatch | None = None,
+    qualname: str | None = None,
+) -> xr.DataArray:
+    """Validate `da` against its declared schema marker(s); return it unchanged.
+
+    Runs each marker's checker under the effective severity, resolved per marker
+    as: the marker's own `on_mismatch` override, else the `on_mismatch` argument,
+    else the policy default (`get_policy().on_mismatch`).  When the master switch
+    is off (`enabled=False`) this is a total no-op.
+
+    Args:
+        da: The DataArray to validate.
+        declared: A single marker or a list of markers declared for `da`.
+        name: The parameter/field name, for error messages.
+        on_mismatch: Per-call default severity, overriding the policy default but
+            overridden by a marker's own `on_mismatch`.
+        qualname: Optional qualified function name, prefixed to messages.
+
+    Returns:
+        `da`, unchanged (validation never mutates).
+
+    Raises:
+        SchemaError: On a mismatch when the effective severity is ``"error"``.
+    """
+    pol = get_policy()
+    if not pol.enabled:
+        return da
+    if on_mismatch is not None:
+        on_mismatch = _validate_on_mismatch(on_mismatch)
+
+    markers = [declared] if isinstance(declared, (Dims, Coords, Dtype)) else declared
+    prefix = f"[{qualname}] {name!r}" if qualname else f"{name!r}"
+    for marker in markers:
+        ok, detail = _CHECKERS[type(marker)](da, marker)
+        severity = marker.on_mismatch or on_mismatch or pol.on_mismatch
+        _report(ok, severity, f"{prefix} {detail}")
+    return da

xarray_annotated/schema/_config.py

@@ -0,0 +1,147 @@
+"""Process-wide validation policy for the structural (schema) domain.
+
+Simpler than the units policy because structural validation never converts — it
+only asserts.  Two axes:
+
+- `enabled` — the **package-wide** master switch (shared with every domain via
+  `xarray_annotated._config`); `False` makes all validation a no-op.
+- `on_mismatch` — what to do when a DataArray's structure does not match a
+  declaration (`"error"` / `"warn"` / `"ignore"`).  This is the *default*; an
+  individual marker may override it (e.g. error on wrong dims but only warn on
+  wrong dtype).
+
+Each axis resolves independently: environment variable, then a process override
+set via `set_policy` / the `policy` context manager, then the module default. The
+shared `_resolve` helper encodes that precedence.
+"""
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Literal
+
+from .._config import (
+    DEFAULT_ENABLED,
+    ENABLED_ENV_VAR,
+    _resolve,
+    _Unset,
+    _UNSET,
+    get_enabled_override,
+    resolve_enabled,
+    set_enabled_override,
+)
+
+# `ENABLED_ENV_VAR` is re-exported (the master switch is shared but callers still
+# reach it as `schema._config.ENABLED_ENV_VAR`).
+__all__ = [
+    "ENABLED_ENV_VAR",
+    "OnMismatch",
+    "Policy",
+    "get_policy",
+    "policy",
+    "set_policy",
+]
+
+OnMismatch = Literal["error", "warn", "ignore"]
+
+_VALID_ON_MISMATCH: frozenset[str] = frozenset({"error", "warn", "ignore"})
+
+#: A structural mismatch is usually a genuine bug (wrong array wired in), so the
+#: default is to raise.  Soften it per-marker (`Dims(..., on_mismatch="warn")`) or
+#: package-wide via the policy.
+DEFAULT_ON_MISMATCH: OnMismatch = "error"
+
+ON_MISMATCH_ENV_VAR = "XARRAY_ANNOTATED_SCHEMA_ON_MISMATCH"
+
+_process_on_mismatch: OnMismatch | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class Policy:
+    """The resolved schema-validation policy — the two axes as a single value.
+
+    Returned by `get_policy`.  Build overrides via `set_policy` or the `policy`
+    context manager rather than constructing this directly.
+
+    Attributes:
+        enabled: Master switch; `False` makes all validation a no-op.
+        on_mismatch: Default behaviour on a structural mismatch
+            (`"error"`, `"warn"`, or `"ignore"`); a marker may override it.
+    """
+
+    enabled: bool = DEFAULT_ENABLED
+    on_mismatch: OnMismatch = DEFAULT_ON_MISMATCH
+
+
+def _validate_on_mismatch(value: str) -> OnMismatch:
+    if value not in _VALID_ON_MISMATCH:
+        raise ValueError(
+            f"Invalid on_mismatch {value!r}. "
+            f"Choose one of {sorted(_VALID_ON_MISMATCH)}."
+        )
+    return value  # type: ignore[return-value]
+
+
+def get_policy() -> Policy:
+    """Resolve the active schema policy (env → process → default, per axis).
+
+    Returns:
+        The resolved `Policy`.
+    """
+    return Policy(
+        enabled=resolve_enabled(),
+        on_mismatch=_resolve(
+            ON_MISMATCH_ENV_VAR,
+            _process_on_mismatch,
+            DEFAULT_ON_MISMATCH,
+            lambda v: _validate_on_mismatch(v.lower()),
+        ),
+    )
+
+
+def set_policy(
+    *,
+    enabled: bool | None | _Unset = _UNSET,
+    on_mismatch: OnMismatch | None | _Unset = _UNSET,
+) -> None:
+    """Set process-wide schema-policy overrides for one or more axes.
+
+    Only the axes you pass are changed.  Pass a value to set it, `None` to
+    *clear* that axis's override (so its env var / default applies again), or
+    omit it to leave it untouched.
+
+    Args:
+        enabled: Override the (package-wide) master switch, or `None` to clear.
+        on_mismatch: Override the on-mismatch axis, or `None` to clear.
+    """
+    global _process_on_mismatch
+    if not isinstance(enabled, _Unset):
+        set_enabled_override(enabled)
+    if not isinstance(on_mismatch, _Unset):
+        _process_on_mismatch = (
+            None if on_mismatch is None else _validate_on_mismatch(on_mismatch)
+        )
+
+
+@contextmanager
+def policy(
+    *,
+    enabled: bool | None | _Unset = _UNSET,
+    on_mismatch: OnMismatch | None | _Unset = _UNSET,
+):
+    """Temporarily override schema-policy axes, restoring them on exit.
+
+    Sets every axis you pass in one go, and restores the previous process
+    overrides afterwards even if an exception is raised.
+
+    Args:
+        enabled: Override the (package-wide) master switch, or `None` to clear.
+        on_mismatch: Override the on-mismatch axis, or `None` to clear.
+    """
+    global _process_on_mismatch
+    saved = (get_enabled_override(), _process_on_mismatch)
+    set_policy(enabled=enabled, on_mismatch=on_mismatch)
+    try:
+        yield
+    finally:
+        enabled_saved, _process_on_mismatch = saved
+        set_enabled_override(enabled_saved)