xarray-annotated 0.5.2__py3-none-any.whl

xarray_annotated/temporal/_config.py

@@ -0,0 +1,182 @@
+"""Process-wide validation policy for the temporal (frequency) domain.
+
+Like the schema policy, temporal validation never converts — it only asserts — but
+it has a *second* behavioural axis, because a frequency declaration can fail in two
+genuinely different ways:
+
+- `on_mismatch` — the array's time axis has a frequency, and it is not the declared
+  one (the declaration is **violated**).
+- `on_uninferable` — no frequency could be determined at all: fewer than three
+  timestamps, or irregular spacing (the declaration could not be **tested**).  This
+  defaults to `"warn"` rather than `"error"`: a short axis is legitimate (a
+  two-timestep test fixture), but silently skipping a contract check deserves noise.
+
+Plus `enabled` — the **package-wide** master switch (shared with every domain via
+`xarray_annotated._config`); `False` makes all validation a no-op.
+
+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 `temporal._config.ENABLED_ENV_VAR`).
+__all__ = [
+    "ENABLED_ENV_VAR",
+    "OnMismatch",
+    "OnUninferable",
+    "Policy",
+    "get_policy",
+    "policy",
+    "set_policy",
+]
+
+OnMismatch = Literal["error", "warn", "ignore"]
+OnUninferable = Literal["error", "warn", "ignore"]
+
+_VALID_SEVERITIES: frozenset[str] = frozenset({"error", "warn", "ignore"})
+
+#: A frequency mismatch is a genuine bug (a weekly series wired into a daily
+#: consumer, or a resample landing on the wrong weekday), so the default is to raise.
+DEFAULT_ON_MISMATCH: OnMismatch = "error"
+
+#: An uninferable axis means the declaration was never *tested*, not that it was
+#: violated — worth noise, but not fatal.
+DEFAULT_ON_UNINFERABLE: OnUninferable = "warn"
+
+ON_MISMATCH_ENV_VAR = "XARRAY_ANNOTATED_TEMPORAL_ON_MISMATCH"
+ON_UNINFERABLE_ENV_VAR = "XARRAY_ANNOTATED_TEMPORAL_ON_UNINFERABLE"
+
+_process_on_mismatch: OnMismatch | None = None
+_process_on_uninferable: OnUninferable | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class Policy:
+    """The resolved temporal-validation policy — the three 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 when the inferred frequency contradicts the
+            declaration (`"error"`, `"warn"`, or `"ignore"`); a marker may override it.
+        on_uninferable: Behaviour when no frequency can be inferred from the time
+            axis — too few points, or irregular spacing.
+    """
+
+    enabled: bool = DEFAULT_ENABLED
+    on_mismatch: OnMismatch = DEFAULT_ON_MISMATCH
+    on_uninferable: OnUninferable = DEFAULT_ON_UNINFERABLE
+
+
+def _validate_severity(value: str, axis: str) -> str:
+    if value not in _VALID_SEVERITIES:
+        raise ValueError(
+            f"Invalid {axis} {value!r}. Choose one of {sorted(_VALID_SEVERITIES)}."
+        )
+    return value
+
+
+def _validate_on_mismatch(value: str) -> OnMismatch:
+    return _validate_severity(value, "on_mismatch")  # type: ignore[return-value]
+
+
+def _validate_on_uninferable(value: str) -> OnUninferable:
+    return _validate_severity(value, "on_uninferable")  # type: ignore[return-value]
+
+
+def get_policy() -> Policy:
+    """Resolve the active temporal 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()),
+        ),
+        on_uninferable=_resolve(
+            ON_UNINFERABLE_ENV_VAR,
+            _process_on_uninferable,
+            DEFAULT_ON_UNINFERABLE,
+            lambda v: _validate_on_uninferable(v.lower()),
+        ),
+    )
+
+
+def set_policy(
+    *,
+    enabled: bool | None | _Unset = _UNSET,
+    on_mismatch: OnMismatch | None | _Unset = _UNSET,
+    on_uninferable: OnUninferable | None | _Unset = _UNSET,
+) -> None:
+    """Set process-wide temporal-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.
+        on_uninferable: Override the on-uninferable axis, or `None` to clear.
+    """
+    global _process_on_mismatch, _process_on_uninferable
+    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)
+        )
+    if not isinstance(on_uninferable, _Unset):
+        _process_on_uninferable = (
+            None if on_uninferable is None else _validate_on_uninferable(on_uninferable)
+        )
+
+
+@contextmanager
+def policy(
+    *,
+    enabled: bool | None | _Unset = _UNSET,
+    on_mismatch: OnMismatch | None | _Unset = _UNSET,
+    on_uninferable: OnUninferable | None | _Unset = _UNSET,
+):
+    """Temporarily override temporal-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.
+        on_uninferable: Override the on-uninferable axis, or `None` to clear.
+    """
+    global _process_on_mismatch, _process_on_uninferable
+    saved = (get_enabled_override(), _process_on_mismatch, _process_on_uninferable)
+    set_policy(enabled=enabled, on_mismatch=on_mismatch, on_uninferable=on_uninferable)
+    try:
+        yield
+    finally:
+        enabled_saved, _process_on_mismatch, _process_on_uninferable = saved
+        set_enabled_override(enabled_saved)

xarray_annotated/temporal/_decorator.py

@@ -0,0 +1,133 @@
+"""Signature-driven frequency-declaration decorator.
+
+`declare_freq` wires the declaration side (`_annotations`) to the checking side
+(`_check`): it reads a function's `Annotated[DataArray, Freq(...)]` hints once, at
+decoration time, then on every call validates the declared inputs and outputs.
+
+Like `schema.declare_schema` — whose scaffold this mirrors — it is the *convenience*
+layer over the public primitives (`freq_from_signature`, `check_freq`,
+`assert_valid_freq`), and it never mutates: inputs and outputs are *checked* and
+passed through unchanged.
+"""
+
+import dataclasses
+import functools
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+import xarray as xr
+
+from ._annotations import Freq, freq_from_signature
+from ._check import assert_valid_freq, check_freq
+from ._config import (
+    OnMismatch,
+    OnUninferable,
+    _validate_on_mismatch,
+    _validate_on_uninferable,
+    get_policy,
+)
+
+
+def declare_freq(
+    func: Callable[..., Any] | None = None,
+    *,
+    on_mismatch: OnMismatch | None = None,
+    on_uninferable: OnUninferable | None = None,
+) -> Callable[..., Any]:
+    """Apply a function's signature-declared frequencies at runtime.
+
+    Reads the decorated function's own type annotations once, via
+    `freq_from_signature`: parameters annotated `Annotated[DataArray, Freq(...)]`
+    declare an input's temporal frequency, and a `TypedDict`/dataclass return (or a
+    bare `Annotated[DataArray, Freq(...)]` return) declares an output's.
+
+    On each call, under the active `Policy` (`get_policy`), the wrapper validates every
+    declared `DataArray` input via `check_freq`, 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, which wins over the policy default), and an axis whose
+    frequency cannot be inferred at all is reported under `on_uninferable`.  When the
+    policy is **disabled** the wrapper is a total no-op.
+
+    Usable bare (`@declare_freq`) or called (`@declare_freq(on_mismatch=...)`).  Every
+    declaration is validated **at decoration time** (`assert_valid_freq`), so a
+    malformed offset string fails fast at import.
+
+    Args:
+        func: The function to decorate (when used bare); `None` when parametrised.
+        on_mismatch: Default mismatch severity for this function, overriding the policy
+            default. `None` (default) resolves per call from `get_policy`.
+        on_uninferable: Severity for an uninferable time axis, overriding the policy
+            default. `None` (default) resolves per call from `get_policy`.
+
+    Returns:
+        A wrapped function that validates declared inputs and outputs.
+    """
+    if on_mismatch is not None:
+        _validate_on_mismatch(on_mismatch)
+    if on_uninferable is not None:
+        _validate_on_uninferable(on_uninferable)
+
+    def decorate(fn: Callable[..., Any]) -> Callable[..., Any]:
+        input_freqs, output_freqs = freq_from_signature(fn)
+        sig = inspect.signature(fn)
+        qualname = getattr(fn, "__qualname__", repr(fn))
+
+        # Fail fast on malformed declarations (independent of policy).
+        for pname, marker in input_freqs.items():
+            assert_valid_freq(marker, f"{qualname} input {pname!r}")
+        _assert_output_declarations(output_freqs, qualname)
+
+        def _check_output(result: Any, eff: OnMismatch, unin: OnUninferable) -> Any:
+            if isinstance(output_freqs, Freq):
+                if isinstance(result, xr.DataArray):
+                    check_freq(result, output_freqs, "return", eff, unin, qualname)
+            elif isinstance(output_freqs, dict):
+                for key, marker in output_freqs.items():
+                    if isinstance(result, dict):
+                        value = result.get(key)
+                    elif dataclasses.is_dataclass(result):
+                        value = getattr(result, key, None)
+                    else:
+                        continue
+                    if isinstance(value, xr.DataArray):
+                        check_freq(value, marker, key, eff, unin, qualname)
+            return result
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            pol = get_policy()
+            if not pol.enabled:
+                # Master switch off: a total no-op (no validation at all).
+                return fn(*args, **kwargs)
+            # Per-call defaults (a marker's own override still wins in check_freq).
+            eff = on_mismatch if on_mismatch is not None else pol.on_mismatch
+            unin = on_uninferable if on_uninferable is not None else pol.on_uninferable
+            if input_freqs:
+                bound = sig.bind_partial(*args, **kwargs)
+                # bind_partial omits parameters left at their default; apply
+                # them so a declared default is validated too.
+                bound.apply_defaults()
+                for pname, val in list(bound.arguments.items()):
+                    marker = input_freqs.get(pname)
+                    if marker is not None and isinstance(val, xr.DataArray):
+                        check_freq(val, marker, pname, eff, unin, qualname)
+            return _check_output(fn(*args, **kwargs), eff, unin)
+
+        wrapper.__annotations__ = dict(fn.__annotations__)
+        return wrapper
+
+    return decorate if func is None else decorate(func)
+
+
+def _assert_output_declarations(
+    output_freqs: dict[str, Freq] | Freq | None,
+    qualname: str,
+) -> None:
+    """Fail-fast-validate declared output markers at decoration time."""
+    if isinstance(output_freqs, Freq):
+        assert_valid_freq(output_freqs, f"{qualname} output")
+    elif isinstance(output_freqs, dict):
+        for key, marker in output_freqs.items():
+            assert_valid_freq(marker, f"{qualname} output {key!r}")

xarray_annotated/units/__init__.py

@@ -0,0 +1,53 @@
+"""Public API for the units domain of xarray-annotated.
+
+Run-time validation of pint/CF units declared via ``typing.Annotated`` on xarray
+DataArrays.
+"""
+
+from ._annotations import Unit, annotated_unit, units_from_signature
+from ._check import (
+    UnitsWarning,
+    apply_output_units,
+    assert_valid_unit,
+    check_units,
+    units_compatible,
+    units_equal,
+)
+from ._config import (
+    OnInexact,
+    OnMissing,
+    OnOutput,
+    Policy,
+    get_policy,
+    policy,
+    set_policy,
+)
+from ._decorator import declare_units
+from ._registry import (
+    get_registry,
+    set_registry,
+    use_cf_units,
+)
+
+__all__ = [
+    "OnInexact",
+    "OnMissing",
+    "OnOutput",
+    "Policy",
+    "Unit",
+    "UnitsWarning",
+    "annotated_unit",
+    "apply_output_units",
+    "assert_valid_unit",
+    "check_units",
+    "declare_units",
+    "get_policy",
+    "get_registry",
+    "policy",
+    "set_policy",
+    "set_registry",
+    "units_compatible",
+    "units_equal",
+    "units_from_signature",
+    "use_cf_units",
+]

xarray_annotated/units/_annotations.py

@@ -0,0 +1,169 @@
+"""Reading declared units off `typing.Annotated` type hints.
+
+Pure `typing` introspection with no dependency on the active registry or
+configuration: this is the "declaration side" of the package, turning
+`Annotated[DataArray, Unit("degC")]` metadata into plain unit strings that the
+checking layer later validates.
+
+The canonical declaration is the self-identifying `Unit` marker.  A bare unit
+string is accepted as a shorthand under a **unit-first** convention: the unit
+must precede any descriptive metadata, e.g.
+`Annotated[DataArray, "m s-1", "description"]`.
+"""
+
+from typing import (
+    Annotated,
+    Any,
+    get_args,
+    get_origin,
+)
+
+from .._annotations import (
+    _is_dataarray_type,
+    unwrap_annotated,
+    walk_signature,
+)
+
+# Re-exported so `units._annotations.unwrap_annotated` stays a stable import.
+__all__ = ["Unit", "annotated_unit", "units_from_signature", "unwrap_annotated"]
+
+
+class Unit:
+    """Typed marker declaring the expected unit of a DataArray.
+
+    This is the **recommended** way to declare a unit.  Inside ``Annotated`` it is
+    self-identifying and independent of metadata ordering, so it composes cleanly
+    with schema markers or other ``Annotated``-based tooling::
+
+        Annotated[xr.DataArray, Unit("degC"), SomeOtherMarker(...)]
+
+    A bare string ``Annotated[xr.DataArray, "degC"]`` is also accepted as a
+    shorthand (the first string in metadata, by convention) and resolves to the
+    same unit.  Prefer the marker whenever the annotation is shared: it removes
+    both the order dependence and the ambiguity a bare string has against a
+    description string.
+
+    The unit string is *not* validated here — parsing is deferred to
+    `assert_valid_unit` at decoration time, exactly as for the bare-string form.
+    """
+
+    __slots__ = ("_unit",)
+
+    def __init__(self, unit: str) -> None:
+        self._unit = unit
+
+    @property
+    def unit(self) -> str:
+        """The declared unit string."""
+        return self._unit
+
+    def __repr__(self) -> str:
+        return f"Unit({self._unit!r})"
+
+    def __str__(self) -> str:
+        # The bare unit string, for readable interpolation (e.g. f"expected {m}").
+        return self._unit
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Unit):
+            return self._unit == other._unit
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        # Defining __eq__ drops the default __hash__; restore it so that an
+        # Annotated[...] hint carrying a Unit marker stays hashable (typing and
+        # downstream tools cache/hash annotations).
+        return hash(self._unit)
+
+
+def annotated_unit(hint: Any) -> str | None:
+    """Return the declared unit carried by an `Annotated` type hint, or `None`.
+
+    The unit is the first `Unit(...)` marker in the `Annotated` metadata, and —
+    falling back — the first `str`.  So both spellings resolve identically:
+
+        `Annotated[DataArray, Unit("degC")]` → `"degC"`
+        `Annotated[DataArray, "degC"]`       → `"degC"`
+
+    The typed `Unit` marker is preferred because it is *self-identifying* and
+    order-independent: it owns its slot regardless of what other metadata
+    (descriptions, other typed markers) shares the annotation, so
+
+        `Annotated[DataArray, "note", Unit("Pa")]` → `"Pa"` (marker wins)
+
+    The bare-string form keeps the extensible *unit-first* convention: the unit
+    may be followed by free-form annotations that are ignored here, so
+
+        `Annotated[DataArray, "m s-1", "z component of velocity"]` → `"m s-1"`
+
+    A description placed before the unit (and with no `Unit` marker present)
+    would be mis-read as the unit — but `assert_valid_unit` rejects it unless the
+    description itself parses as a valid unit, so the failure is loud.  Non-string
+    metadata (ints, unrelated markers) is skipped regardless of position.
+
+    The metadata is only interpreted as a unit when the annotated base type
+    is a `DataArray` (the only type that carries units); a `Unit` marker or
+    descriptive string on a non-`DataArray` parameter (e.g.
+    `Annotated[bool, "toggles X"]`) is *not* a unit and yields `None`.
+    Non-`Annotated` hints, or `Annotated` hints whose metadata holds neither a
+    `Unit` marker nor a string, also return `None`.
+
+    Args:
+        hint: A type hint, typically from `get_type_hints(..., include_extras=True)`.
+
+    Returns:
+        The declared unit (first `Unit` marker, else first `str`) if the base
+        type is a `DataArray`; `None` otherwise.
+
+    Examples:
+        >>> from typing import Annotated
+        >>> import xarray as xr
+        >>> from xarray_annotated.units._annotations import Unit, annotated_unit
+        >>> annotated_unit(Annotated[xr.DataArray, Unit("degC")])
+        'degC'
+        >>> annotated_unit(Annotated[xr.DataArray, "degC"])  # shorthand
+        'degC'
+        >>> annotated_unit(Annotated[bool, "toggles X"]) is None
+        True
+    """
+    if get_origin(hint) is not Annotated:
+        return None
+    # get_args(Annotated[T, m1, m2, ...]) == (T, m1, m2, ...); skip the base type.
+    args = get_args(hint)
+    if not _is_dataarray_type(args[0]):
+        return None
+    metadata = args[1:]
+    # The typed marker wins over a bare string, regardless of relative order.
+    for meta in metadata:
+        if isinstance(meta, Unit):
+            return meta.unit
+    for meta in metadata:
+        if isinstance(meta, str):
+            return meta
+    return None
+
+
+def units_from_signature(
+    func: object,
+) -> tuple[dict[str, str], dict[str, str] | str | None]:
+    """Extract declared units from a function's type annotations.
+
+    Reads `get_type_hints(func, include_extras=True)` and interprets a
+    `Unit(...)` marker (or its bare-string shorthand) in the `Annotated`
+    metadata as a unit declaration.
+
+    Args:
+        func: A callable whose type hints carry `Annotated[DataArray, Unit(...)]`
+            metadata (or the bare-string shorthand).
+
+    Returns:
+        A `(input_units, output_units)` pair:
+
+        * `input_units` — `dict[str, str]` mapping parameter names to their
+            declared unit strings.  Only parameters whose hint carries a unit
+            declaration contribute.
+        * `output_units` — a `dict[str, str]` (per-field units) if the
+            return hint is a `TypedDict` or dataclass; a bare `str` if the return
+            is a single `Annotated[DataArray, Unit(...)]`; `None` otherwise.
+    """
+    return walk_signature(func, annotated_unit)