xarray-annotated 0.5.2__py3-none-any.whl

xarray_annotated/__init__.py

@@ -0,0 +1,50 @@
+"""xarray-annotated: validate xarray DataArray properties declared via ``typing.Annotated``.
+
+Annotation is the unifying technology: a declared property is read off an
+``Annotated[DataArray, ...]`` hint and validated at runtime.  Each kind of
+property is a domain subpackage, imported explicitly:
+
+* ``xarray_annotated.units`` — physical units (pint / CF), the mature domain.
+* ``xarray_annotated.schema`` — structural properties (dims, coords, dtype);
+  validate-only (never mutates).
+* ``xarray_annotated.temporal`` — the frequency (and phase) of a time axis;
+  validate-only.
+
+The top level is deliberately thin — nothing domain-specific is re-exported
+here, so the domains never collide in a shared namespace::
+
+    from xarray_annotated import schema, temporal, units
+    from xarray_annotated.units import declare_units, check_units
+    from xarray_annotated.schema import declare_schema, Dims, Dtype
+    from xarray_annotated.temporal import declare_freq, Freq
+
+The only names surfaced at the top level are domain-*agnostic* helpers that
+belong to no single domain: the shared declaration *writer* ``annotate`` and its
+inverse *reader* ``declarations_from_signature`` (which reads every facet declared
+on a signature into one uniform ``Declared`` value), plus the ``Annotated``
+introspection kernel (``unwrap_annotated``, ``walk_signature``).
+``walk_signature`` is the shared driver behind every domain's reader, so a
+third-party facet author can use it to build their own ``*_from_signature``
+reader.  No domain-specific name is re-exported here.
+"""
+
+from importlib.metadata import version as _version
+
+from . import schema, temporal, units
+from ._annotations import unwrap_annotated, walk_signature
+from ._reader import Declared, declarations_from_signature
+from ._writer import annotate
+
+__version__ = _version("xarray-annotated")
+
+__all__ = [
+    "Declared",
+    "__version__",
+    "annotate",
+    "declarations_from_signature",
+    "schema",
+    "temporal",
+    "units",
+    "unwrap_annotated",
+    "walk_signature",
+]

xarray_annotated/_annotations.py

@@ -0,0 +1,119 @@
+"""Shared `typing.Annotated` introspection for xarray-annotated.
+
+Pure `typing` machinery with no dependency on any domain's registry or config —
+the "declaration side" common to every domain. A domain supplies an *extractor*
+that turns one `Annotated[...]` hint into its own declaration payload (a unit
+string, a list of schema markers, …); `walk_signature` drives that extractor over
+a whole function signature, handling the `TypedDict` / dataclass / single-value
+return shapes once, for everyone.
+"""
+
+import dataclasses
+import types
+from collections.abc import Callable
+from typing import (
+    Annotated,
+    Any,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+    is_typeddict,
+)
+
+import xarray as xr
+
+_T = TypeVar("_T")
+
+
+def unwrap_annotated(hint: Any) -> Any:
+    """Return the underlying type of an `Annotated` hint, else the hint itself.
+
+    `Annotated[DataArray, ...]` → `DataArray`; a non-`Annotated` hint is returned
+    unchanged.  Lets type comparisons (e.g. `t is DataArray`) see through the
+    metadata that declarations attach to parameters.
+
+    Args:
+        hint: A type hint, possibly `Annotated`.
+
+    Returns:
+        The base type if `hint` is `Annotated`; otherwise `hint` unchanged.
+    """
+    return get_args(hint)[0] if get_origin(hint) is Annotated else hint
+
+
+def _is_dataarray_type(tp: Any) -> bool:
+    """Return whether `tp` is `DataArray` (possibly wrapped in a Union).
+
+    Accepts a bare `DataArray` as well as `DataArray | None` /
+    `Optional[DataArray]` (and any other union that includes `DataArray`), so an
+    optional DataArray parameter can still carry a declaration.  Anything else
+    (scalars, `str`, `bool`, `Dataset`, …) is not a declaration-bearing type.
+
+    Args:
+        tp: A type annotation to inspect.
+
+    Returns:
+        `True` if `tp` is or contains `DataArray`.
+    """
+    if tp is xr.DataArray:
+        return True
+    if get_origin(tp) in (Union, types.UnionType):
+        return any(_is_dataarray_type(arg) for arg in get_args(tp))
+    return False
+
+
+def walk_signature(
+    func: object,
+    extract: Callable[[Any], _T | None],
+) -> tuple[dict[str, _T], dict[str, _T] | _T | None]:
+    """Extract per-parameter and return declarations from a function's hints.
+
+    Reads `get_type_hints(func, include_extras=True)` and maps each hint through
+    `extract`, which returns a domain's declaration payload for that hint or
+    `None` when the hint carries no declaration.  The `is not None` filter is the
+    single "no declaration" rule shared by every domain.
+
+    Args:
+        func: A callable whose parameters/return may carry `Annotated` metadata.
+        extract: Turns one hint into a declaration payload, or `None`.  For units
+            this is `annotated_unit` (payload `str`); for schema `annotated_schema`
+            (payload `list[marker]`).
+
+    Returns:
+        An `(inputs, output)` pair:
+
+        * `inputs` — `dict[str, T]` mapping each declared parameter to its payload.
+        * `output` — a `dict[str, T]` (per-field) if the return hint is a
+            `TypedDict` or dataclass; a single `T` if the return is one declared
+            `Annotated[DataArray, ...]`; `None` otherwise.
+    """
+    hints = get_type_hints(func, include_extras=True)
+    ret = hints.pop("return", None)
+
+    inputs = {
+        name: payload
+        for name, hint in hints.items()
+        if (payload := extract(hint)) is not None
+    }
+
+    output: dict[str, _T] | _T | None
+    if is_typeddict(ret):
+        ret_hints = get_type_hints(ret, include_extras=True)
+        output = {
+            name: payload
+            for name, hint in ret_hints.items()
+            if (payload := extract(hint)) is not None
+        }
+    elif dataclasses.is_dataclass(ret):
+        ret_hints = get_type_hints(ret, include_extras=True)
+        output = {
+            f.name: payload
+            for f in dataclasses.fields(ret)
+            if (payload := extract(ret_hints.get(f.name))) is not None
+        }
+    else:
+        output = extract(ret)
+
+    return inputs, output

xarray_annotated/_config.py

@@ -0,0 +1,101 @@
+"""Shared configuration kernel for xarray-annotated.
+
+Machinery common to every annotation-validation domain (currently `units` and
+`schema`), so the domains never drift on how a policy axis resolves:
+
+- `_resolve` — the shared precedence for one axis: environment variable (parsed),
+  then a process override, then a module default.
+- the **`enabled`** master switch — package-wide (it gates *every* domain, not just
+  one), so it lives here rather than in any single domain. Each domain's `Policy`
+  carries an `enabled` field but resolves it through `resolve_enabled` / the
+  override accessors below, so toggling it in one domain toggles it everywhere.
+- `_Unset` / `_UNSET` — the sentinel that lets a `set_policy` keyword distinguish
+  "argument omitted" from "explicitly set to `None`" (which *clears* an override).
+
+Domain-specific axes (units' `on_missing` / `on_inexact`, schema's `on_mismatch`)
+stay in the respective domain's `_config.py`; only the genuinely shared pieces
+live here.
+"""
+
+import os
+from collections.abc import Callable
+from typing import TypeVar
+
+#: Master switch default. Validation is on unless explicitly disabled.
+DEFAULT_ENABLED: bool = True
+
+#: Master switch is package-wide (gates every annotation-validation domain), so it
+#: is deliberately un-namespaced; each domain's behavioural knobs are namespaced.
+ENABLED_ENV_VAR = "XARRAY_ANNOTATED_ENABLED"
+
+#: String values (lower-cased) accepted for the boolean `enabled` env var.
+_TRUTHY: frozenset[str] = frozenset({"1", "true", "yes", "on"})
+_FALSEY: frozenset[str] = frozenset({"0", "false", "no", "off"})
+
+#: Process-wide override for the master switch, shared across all domains.
+_process_enabled: bool | None = None
+
+_T = TypeVar("_T")
+
+
+class _Unset:
+    """Sentinel distinguishing 'argument omitted' from 'explicitly set to None'."""
+
+    def __repr__(self) -> str:
+        return "<unset>"
+
+
+_UNSET = _Unset()
+
+
+def _resolve(
+    env_var: str,
+    process_value: _T | None,
+    default: _T,
+    parse: Callable[[str], _T],
+) -> _T:
+    """Resolve one axis: env var (parsed) → process override → default."""
+    env = os.environ.get(env_var)
+    if env:
+        return parse(env)
+    if process_value is not None:
+        return process_value
+    return default
+
+
+def _parse_enabled_env(value: str) -> bool:
+    lowered = value.lower()
+    if lowered in _TRUTHY:
+        return True
+    if lowered in _FALSEY:
+        return False
+    raise ValueError(
+        f"Invalid {ENABLED_ENV_VAR} value {value!r}. "
+        f"Use one of {sorted(_TRUTHY | _FALSEY)}."
+    )
+
+
+def resolve_enabled() -> bool:
+    """Resolve the package-wide master switch (env → process override → default)."""
+    return _resolve(
+        ENABLED_ENV_VAR, _process_enabled, DEFAULT_ENABLED, _parse_enabled_env
+    )
+
+
+def get_enabled_override() -> bool | None:
+    """Return the current process override for the master switch (`None` if unset).
+
+    Used by each domain's `policy` context manager to save/restore the shared
+    switch around a temporary override.
+    """
+    return _process_enabled
+
+
+def set_enabled_override(value: bool | None) -> None:
+    """Set (or, with `None`, clear) the process override for the master switch.
+
+    Shared by every domain's `set_policy`, so disabling validation in one domain
+    disables it package-wide.
+    """
+    global _process_enabled
+    _process_enabled = None if value is None else bool(value)

xarray_annotated/_reader.py

@@ -0,0 +1,110 @@
+"""Reading *all* declared facets off a signature — the unified reader.
+
+Every domain ships its own reader (`units_from_signature`, `schema_from_signature`,
+`freq_from_signature`) returning a domain-shaped payload; this is the one
+cross-domain reader that collapses them into a single uniform value.  Where the
+writer `annotate` turns facet values into an `Annotated` hint,
+`declarations_from_signature` turns a whole signature back into a facet value per
+declared DataArray — it is the read-side counterpart to `annotate`, and their
+round-trip is exact.
+
+The unified value `Declared` is deliberately *homogeneous*: every slot holds a
+marker (`Unit` / `Dims` / `Dtype` / `Coords` / `Freq`) or `None`.  Markers are the
+package's public currency — the schema and temporal markers are load-bearing (they
+carry strictness that has no bare-string form) — so a facet-generic consumer wants
+one marker per slot rather than a mix of strings and objects.  A bare-string unit
+shorthand is normalised to a `Unit`, so the shape does not depend on how the unit
+was spelled.
+
+Like `annotate`, this is cross-domain — it assembles every domain's markers into one
+value — so it lives at the package root rather than in any one domain.
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+from ._annotations import walk_signature
+from .schema import Coords, Dims, Dtype
+from .schema._annotations import annotated_schema
+from .temporal import Freq
+from .temporal._annotations import annotated_freq
+from .units import Unit
+from .units._annotations import annotated_unit
+
+__all__ = ["Declared", "declarations_from_signature"]
+
+
+@dataclass(frozen=True, slots=True)
+class Declared:
+    """Every facet declared on one DataArray hint, one marker (or `None`) per slot.
+
+    The uniform value produced by `declarations_from_signature`: a homogeneous
+    bag of markers, so facet-generic code can treat each slot identically.  Each
+    slot holds the corresponding marker or `None` when that facet was not
+    declared; the unit is always a `Unit` (its bare-string shorthand is
+    normalised on read), so `.unit.unit` recovers the string.
+
+    The field order (`unit`, `dims`, `dtype`, `coords`, `freq`) matches `annotate`'s
+    keyword arguments, so `annotate(unit=d.unit, dims=d.dims, dtype=d.dtype,
+    coords=d.coords, freq=d.freq)` rebuilds the hint this was read from.
+    """
+
+    unit: Unit | None = None
+    dims: Dims | None = None
+    dtype: Dtype | None = None
+    coords: Coords | None = None
+    freq: Freq | None = None
+
+
+def _annotated_declaration(hint: Any) -> Declared | None:
+    """Combine the per-domain extractors into one `Declared`, or `None`.
+
+    Delegates to the existing `annotated_unit`, `annotated_schema` and
+    `annotated_freq` so the "is this an `Annotated[DataArray, ...]` hint?" logic is
+    not duplicated, then folds their payloads into a single homogeneous value.
+    Returns `None` when the hint declares no facet, which is `walk_signature`'s "no
+    declaration" signal.
+    """
+    unit = annotated_unit(hint)
+    markers = annotated_schema(hint) or ()
+    dims = next((m for m in markers if isinstance(m, Dims)), None)
+    dtype = next((m for m in markers if isinstance(m, Dtype)), None)
+    coords = next((m for m in markers if isinstance(m, Coords)), None)
+    freq = annotated_freq(hint)
+    if all(facet is None for facet in (unit, dims, dtype, coords, freq)):
+        return None
+    return Declared(
+        unit=Unit(unit) if unit is not None else None,
+        dims=dims,
+        dtype=dtype,
+        coords=coords,
+        freq=freq,
+    )
+
+
+def declarations_from_signature(
+    func: object,
+) -> tuple[dict[str, Declared], dict[str, Declared] | Declared | None]:
+    """Read every declared facet off a function's signature in one uniform shape.
+
+    The all-facets counterpart to `units_from_signature` / `schema_from_signature`:
+    a single reader whose payload is a `Declared` carrying the unit, dims, dtype,
+    and coords declared on a DataArray hint at once.  A consumer that would
+    otherwise call both per-domain readers and merge their differently-shaped
+    results can call this once and route a uniform value instead.
+
+    Args:
+        func: A callable whose parameters/return may carry
+            `Annotated[DataArray, ...]` metadata.
+
+    Returns:
+        An `(inputs, output)` pair, mirroring `walk_signature` with a `Declared`
+        payload:
+
+        * `inputs` — `dict[str, Declared]` mapping each declared parameter to its
+            facets.  Parameters carrying no declaration are omitted.
+        * `output` — a `dict[str, Declared]` (per-field) for a `TypedDict` or
+            dataclass return, a single `Declared` for one declared DataArray
+            return, or `None` when the return carries no declaration.
+    """
+    return walk_signature(func, _annotated_declaration)

xarray_annotated/_writer.py

@@ -0,0 +1,99 @@
+"""Writing declarations *onto* an annotation — the inverse of the readers.
+
+Every domain ships a reader (`Annotated → declaration`: `units_from_signature`,
+`schema_from_signature`, `freq_from_signature`); this is the one shared *writer*
+(`declaration → Annotated`).  A tool that builds functions dynamically (codegen,
+graph nodes) needs to stamp a declared contract onto a generated signature;
+`annotate` returns a real `Annotated` object it can assign to
+`fn.__annotations__["return"]`, so the ordinary `declare_units` / `declare_schema` /
+`declare_freq` decorators then read it straight back.
+
+It is deliberately cross-domain — it assembles the units `Unit` marker, the schema
+`Dims` / `Dtype` / `Coords` markers, and the temporal `Freq` marker — so it lives at
+the package root rather than in any one domain, next to the domain-agnostic
+annotation helpers.
+"""
+
+from collections.abc import Iterable
+from typing import Annotated, Any
+
+import xarray as xr
+
+from .schema import Coords, Dims, Dtype
+from .temporal import Freq
+from .units import Unit
+
+__all__ = ["annotate"]
+
+
+def annotate(
+    base: Any = xr.DataArray,
+    *,
+    unit: str | Unit | None = None,
+    dims: Iterable[str] | Dims | None = None,
+    dtype: str | Dtype | None = None,
+    coords: Iterable[str] | Coords | None = None,
+    freq: str | Freq | None = None,
+) -> Any:
+    """Build an `Annotated[base, <markers>]` hint from declared facet values.
+
+    The inverse of the `*_from_signature` readers: given facet values it returns a
+    real `Annotated` object carrying the corresponding markers, in a fixed order
+    (unit, dims, dtype, coords, freq).  Assign it to a function's return/parameter
+    annotation and the `declare_units` / `declare_schema` / `declare_freq` decorators
+    read it back exactly as they would a hand-written one.
+
+    Each facet accepts either a raw value or an already-built marker, so a caller
+    holding a mix (e.g. a `Unit` object but bare dim-name tuples) can pass both
+    without unwrapping:
+
+        * `unit`   — a unit string (`"Pa"`) or a `Unit`.
+        * `dims`   — an iterable of dim names (`("time", "x")`) or a `Dims`.
+        * `dtype`  — a dtype string (`"float64"`) or a `Dtype`.
+        * `coords` — an iterable of coord names or a `Coords`.
+        * `freq`   — an offset string (`"7D"`) or a `Freq`.
+
+    A facet left as `None` contributes no marker.  When no facet is given, `base`
+    is returned unchanged (no `Annotated` wrapper), so `annotate()` is a safe
+    no-op default.
+
+    Note the `freq` string is a *writer* convenience only: there is no bare-string
+    shorthand on the read side, so a frequency must be spelled as a `Freq` marker in
+    a hand-written annotation.
+
+    Args:
+        base: The base type to annotate (default `xarray.DataArray`).
+        unit: Declared unit, or `None`.
+        dims: Declared dimensions, or `None`.
+        dtype: Declared dtype, or `None`.
+        coords: Declared coordinates, or `None`.
+        freq: Declared temporal frequency, or `None`.
+
+    Returns:
+        `Annotated[base, <markers>]` if any facet was given; otherwise `base`.
+
+    Examples:
+        >>> from typing import Annotated, get_args, get_origin
+        >>> import xarray as xr
+        >>> from xarray_annotated import annotate
+        >>> from xarray_annotated.units import Unit
+        >>> hint = annotate(unit="Pa", dims=("time", "x"), dtype="float64")
+        >>> get_origin(hint) is Annotated
+        True
+        >>> Unit("Pa") in get_args(hint)
+        True
+        >>> annotate() is xr.DataArray
+        True
+    """
+    markers: list[Any] = []
+    if unit is not None:
+        markers.append(unit if isinstance(unit, Unit) else Unit(unit))
+    if dims is not None:
+        markers.append(dims if isinstance(dims, Dims) else Dims(*dims))
+    if dtype is not None:
+        markers.append(dtype if isinstance(dtype, Dtype) else Dtype(dtype))
+    if coords is not None:
+        markers.append(coords if isinstance(coords, Coords) else Coords(*coords))
+    if freq is not None:
+        markers.append(freq if isinstance(freq, Freq) else Freq(freq))
+    return Annotated[(base, *markers)] if markers else base

xarray_annotated/schema/__init__.py

@@ -0,0 +1,62 @@
+"""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:
+        ...
+"""
+
+from ._annotations import (
+    Coords,
+    Dims,
+    Dtype,
+    annotated_schema,
+    schema_from_signature,
+)
+from ._check import (
+    SchemaError,
+    SchemaWarning,
+    assert_valid_schema,
+    check_schema,
+    dims_compatible,
+    dtype_compatible,
+)
+from ._config import (
+    OnMismatch,
+    Policy,
+    get_policy,
+    policy,
+    set_policy,
+)
+from ._decorator import declare_schema
+
+__all__ = [
+    "Coords",
+    "Dims",
+    "Dtype",
+    "OnMismatch",
+    "Policy",
+    "SchemaError",
+    "SchemaWarning",
+    "annotated_schema",
+    "assert_valid_schema",
+    "check_schema",
+    "declare_schema",
+    "dims_compatible",
+    "dtype_compatible",
+    "get_policy",
+    "policy",
+    "schema_from_signature",
+    "set_policy",
+]