Skip to content

xarray_annotated.schema

xarray_annotated.schema

Structural-property domain for xarray-annotated.

Declares and validates a DataArray’s structural properties — the ones every DataArray possesses regardless of physical units: its dimensions (Dims), coordinates (Coords), and dtype (Dtype). The structural counterpart to xarray_annotated.units, but simpler: structural validation only asserts (it never converts or mutates).

Declare several at once inside Annotated and apply them with declare_schema::

from typing import Annotated
import xarray as xr
from xarray_annotated.schema import declare_schema, Dims, Dtype

@declare_schema
def f(
    x: Annotated[xr.DataArray, Dims("time", "x"), Dtype("float64")],
) -> xr.DataArray: ...

Coords 

Coords(*names: str, on_mismatch: OnMismatch | None = None)

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 property 

names: tuple[str, ...]

The declared coordinate names.

on_mismatch property 

on_mismatch: OnMismatch | None

Per-marker severity override, or None to use the policy default.

Dims 

Dims(*names: str, ordered: bool = False, on_mismatch: OnMismatch | None = None)

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 property 

names: tuple[str, ...]

The declared dimension names, in the order given.

on_mismatch property 

on_mismatch: OnMismatch | None

Per-marker severity override, or None to use the policy default.

ordered property 

ordered: bool

Whether dim order (not just the set) must match.

Dtype 

Dtype(dtype: str, *, exact: bool = False, on_mismatch: OnMismatch | None = None)

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 property 

dtype: str

The declared dtype string.

exact property 

exact: bool

Whether the exact dtype (not just its kind) must match.

on_mismatch property 

on_mismatch: OnMismatch | None

Per-marker severity override, or None to use the policy default.

Policy dataclass 

Policy(enabled: bool = DEFAULT_ENABLED, on_mismatch: OnMismatch = DEFAULT_ON_MISMATCH)

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 (bool) –

    Master switch; False makes all validation a no-op.

  • on_mismatch (OnMismatch) –

    Default behaviour on a structural mismatch ("error", "warn", or "ignore"); a marker may override it.

SchemaError

Bases: 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.

SchemaWarning

Bases: UserWarning

Warning issued for a structural mismatch under on_mismatch="warn".

annotated_schema 

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).

Parameters:

  • hint (Any) –

    A type hint, typically from get_type_hints(..., include_extras=True).

Returns:

  • list[SchemaMarker] | None –

    A non-empty list of markers in annotation order, or None.

assert_valid_schema 

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.

Parameters:

  • marker (SchemaMarker) –

    The declared Dims, Coords, or Dtype marker.

  • context (str) –

    A label for the error message (e.g. "f input 'x'").

Raises:

check_schema 

check_schema(
    da: DataArray,
    declared: SchemaMarker | list[SchemaMarker],
    name: str,
    on_mismatch: OnMismatch | None = None,
    qualname: str | None = None,
) -> 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.

Parameters:

  • da (DataArray) –

    The DataArray to validate.

  • declared (SchemaMarker | list[SchemaMarker]) –

    A single marker or a list of markers declared for da.

  • name (str) –

    The parameter/field name, for error messages.

  • on_mismatch (OnMismatch | None, default: None ) –

    Per-call default severity, overriding the policy default but overridden by a marker’s own on_mismatch.

  • qualname (str | None, default: None ) –

    Optional qualified function name, prefixed to messages.

Returns:

  • DataArray –

    da, unchanged (validation never mutates).

Raises:

  • SchemaError –

    On a mismatch when the effective severity is "error".

declare_schema 

declare_schema(
    func: Callable[..., Any] | None = None, *, on_mismatch: OnMismatch | None = None
) -> Callable[..., Any]

Apply a function’s signature-declared structure at runtime.

Reads the decorated function’s own type annotations once, via schema_from_signature: parameters annotated Annotated[DataArray, <markers>] declare input structure, and a TypedDict/dataclass return (or a bare Annotated[DataArray, <markers>] return) declares output structure.

On each call, under the active Policy (get_policy), the wrapper validates every declared DataArray input via check_schema, runs the wrapped function, then validates every declared output the same way. Nothing is mutated; a mismatch raises/warns/ignores per the effective on_mismatch (a marker’s own override wins over this decorator’s on_mismatch, which wins over the policy default). When the policy is disabled the wrapper is a total no-op.

Usable bare (@declare_schema) or called (@declare_schema(on_mismatch=...)). Every marker declaration is validated at decoration time (assert_valid_schema), so a malformed declaration fails fast at import.

Parameters:

  • func (Callable[..., Any] | None, default: None ) –

    The function to decorate (when used bare); None when parametrised.

  • on_mismatch (OnMismatch | None, default: None ) –

    Default severity for this function, overriding the policy default. None (default) resolves per call from get_policy.

Returns:

  • Callable[..., Any] –

    A wrapped function that validates declared inputs and outputs.

dims_compatible 

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.

Parameters:

  • a (Dims) –

    A Dims marker.

  • b (Dims) –

    Another Dims marker.

Returns:

  • bool –

    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

dtype_compatible 

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.

Parameters:

  • a (Dtype) –

    A Dtype marker.

  • b (Dtype) –

    Another Dtype marker.

Returns:

  • bool –

    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

get_policy 

get_policy() -> Policy

Resolve the active schema policy (env → process → default, per axis).

Returns:

  • Policy –

    The resolved Policy.

policy 

policy(
    *,
    enabled: bool | _Unset | None = _UNSET,
    on_mismatch: OnMismatch | _Unset | None = _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.

Parameters:

  • enabled (bool | _Unset | None, default: _UNSET ) –

    Override the (package-wide) master switch, or None to clear.

  • on_mismatch (OnMismatch | _Unset | None, default: _UNSET ) –

    Override the on-mismatch axis, or None to clear.

schema_from_signature 

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.

Parameters:

  • func (object) –

    A callable whose hints carry Annotated[DataArray, <markers>].

Returns:

  • dict[str, list[SchemaMarker]] –

    An (inputs, output) pair, where each declaration is the list of markers

  • dict[str, list[SchemaMarker]] | list[SchemaMarker] | None –

    on that parameter/field; output is a per-field dict for a TypedDict or

  • tuple[dict[str, list[SchemaMarker]], dict[str, list[SchemaMarker]] | list[SchemaMarker] | None] –

    dataclass return, a single list for one declared DataArray return, or

  • tuple[dict[str, list[SchemaMarker]], dict[str, list[SchemaMarker]] | list[SchemaMarker] | None] –

    None.

set_policy 

set_policy(
    *,
    enabled: bool | _Unset | None = _UNSET,
    on_mismatch: OnMismatch | _Unset | None = _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.

Parameters:

  • enabled (bool | _Unset | None, default: _UNSET ) –

    Override the (package-wide) master switch, or None to clear.

  • on_mismatch (OnMismatch | _Unset | None, default: _UNSET ) –

    Override the on-mismatch axis, or None to clear.