xarray-annotated 0.5.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

xarray_annotated/__init__.py

@@ -1,50 +1,60 @@
-"""xarray-annotated: validate xarray DataArray properties declared via ``typing.Annotated``.
+"""xarray-annotated: validate 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
+``Annotated[DataArray, ...]`` hint and enforced at runtime.  Each kind of
 property is a domain subpackage, imported explicitly:
 
-* ``xarray_annotated.units`` — physical units (pint / CF), the mature domain.
+* ``xarray_annotated.units`` — physical units (pint / CF).
 * ``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.
+One decorator, ``enforce``, applies every facet declared on a signature, units first:
+
+    import xarray_annotated as xra
+    from xarray_annotated import Dims, Freq, Unit
+
+    @xra.enforce
+    def f(p: Annotated[xr.DataArray, Unit("Pa"), Dims("time"), Freq("D")]): ...
+
+The top level is deliberately thin — the domains' checker names stay in their
+subpackages.  The exceptions are the markers (``Coords``, ``Dims``, ``Dtype``,
+``Freq``, ``Unit``), which are unique across domains and are combined freely in one
+hint, and the domain-*agnostic*
+names: ``enforce``; the ``enabled`` switch; the shared
+declaration *writer* ``annotate`` and its inverse *readers*
+``declaration_from_hint`` and ``declarations_from_signature`` (which read every facet
+declared on a hint or a signature into one uniform ``Declared`` value).
 """
 
 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 ._config import enabled
+from ._enforce import enforce
+from ._reader import Declared, declaration_from_hint, declarations_from_signature
 from ._writer import annotate
+from .schema import Coords, Dims, Dtype
+from .temporal import Freq
+from .units import Unit
 
 __version__ = _version("xarray-annotated")
 
 __all__ = [
+    "Coords",
     "Declared",
+    "Dims",
+    "Dtype",
+    "Freq",
+    "Unit",
     "__version__",
     "annotate",
+    "declaration_from_hint",
     "declarations_from_signature",
+    "enabled",
+    "enforce",
     "schema",
     "temporal",
     "units",
-    "unwrap_annotated",
-    "walk_signature",
 ]

xarray_annotated/_annotations.py

@@ -1,20 +1,19 @@
 """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.
+the "declaration side" common to every domain. An *extractor* turns one
+`Annotated[...]` hint into a declaration payload;
+`_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 collections.abc import Callable, Iterable
 from typing import (
     Annotated,
     Any,
-    TypeVar,
+    TypeAliasType,
     Union,
     get_args,
     get_origin,
@@ -24,96 +23,157 @@ from typing import (
 
 import xarray as xr
 
-_T = TypeVar("_T")
 
+def _single_marker[M](metadata: Iterable[Any], marker_type: type[M]) -> M | None:
+    """Return the one `marker_type` marker in `metadata`, or `None` if absent.
 
-def unwrap_annotated(hint: Any) -> Any:
-    """Return the underlying type of an `Annotated` hint, else the hint itself.
+    A hint declares each facet at most once.  Two markers of one type are either
+    contradictory (`Dims("x")` with `Dims("y")`) or redundant, so a second one is a
+    malformed declaration rather than something to resolve by position.
 
-    `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.
+    Raises:
+        ValueError: If `metadata` holds more than one `marker_type` marker.
     """
-    return get_args(hint)[0] if get_origin(hint) is Annotated else hint
+    found = [m for m in metadata if isinstance(m, marker_type)]
+    if len(found) > 1:
+        raise ValueError(
+            f"one Annotated hint declares {len(found)} {marker_type.__name__} "
+            f"markers ({', '.join(map(repr, found))}); declare at most one"
+        )
+    return found[0] if found else None
 
 
 def _is_dataarray_type(tp: Any) -> bool:
-    """Return whether `tp` is `DataArray` (possibly wrapped in a Union).
+    """Return whether `tp` is `DataArray` (possibly wrapped or 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.
 
+    Also accepts any *array-type wrapper* that exposes the underlying array class
+    as an `array_type` attribute — notably jaxtyping's `Float[DataArray, "t x"]`,
+    which occupies the base-type slot of an `Annotated` hint rather than competing
+    with it.  This is a duck-typed check: jaxtyping is never imported, and any
+    other library following the same convention composes for free.  See
+    `docs/guides/jaxtyping.md`.
+
     Args:
         tp: A type annotation to inspect.
 
     Returns:
-        `True` if `tp` is or contains `DataArray`.
+        `True` if `tp` is, wraps, or contains `DataArray`.
     """
     if tp is xr.DataArray:
         return True
+    inner = getattr(tp, "array_type", None)
+    if isinstance(inner, type) and issubclass(inner, 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(
+def _strip_return_hint(
+    hint: Any, path: frozenset[Any] = frozenset()
+) -> tuple[Any, bool]:
+    """Strip `type` aliases, `Annotated` and `| None` off a return hint.
+
+    An alias that refers to itself is left in place: it resolves to nothing.  The reader
+    raises on one, except a generic one, as in `type G[T] = G[T] | None`, whose `G[int]`
+    it does not enter.
+
+    Returns:
+        What remains, and whether a `| None` was stripped on the way to it.
+    """
+    origin = get_origin(hint)
+    alias = next((a for a in (hint, origin) if isinstance(a, TypeAliasType)), None)
+    if alias is not None:
+        if alias in path:
+            return hint, False
+        return _strip_return_hint(alias.__value__, path | {alias})
+    if origin is Annotated:
+        return _strip_return_hint(get_args(hint)[0], path)
+    if origin in (Union, types.UnionType):
+        members = [arg for arg in get_args(hint) if arg is not types.NoneType]
+        if len(members) == 1:
+            return _strip_return_hint(members[0], path)[0], True
+    return hint, False
+
+
+def _record_return(hint: Any) -> tuple[Any, bool]:
+    """Read a return hint's `TypedDict` or dataclass, and whether it may be `None`.
+
+    Returns:
+        The class, or `None` if the hint resolves to no such class, and whether the
+        hint allows `None`.
+    """
+    base, optional = _strip_return_hint(hint)
+    record = get_origin(base) or base
+    is_record = is_typeddict(record) or dataclasses.is_dataclass(record)
+    return (record if isinstance(record, type) and is_record else None), optional
+
+
+def _walk_signature[T](
     func: object,
-    extract: Callable[[Any], _T | None],
-) -> tuple[dict[str, _T], dict[str, _T] | _T | None]:
+    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.
+    `extract`, which returns a declaration payload for that hint or `None` when
+    the hint carries no declaration.  `extract` also sees a whole return hint before
+    any of its fields, so it can refuse what it finds there.
 
     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]`).
+        extract: Turns one hint into a declaration payload, or `None`.  For every
+            built-in facet at once this is `declaration_from_hint` (payload
+            `Declared`).
 
     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.
+            `TypedDict` or dataclass, read through a `type` alias, `Annotated`,
+            `| None` and generic parameters; a single `T` if the return is one
+            declared `Annotated[DataArray, ...]`; `None` otherwise.
     """
+    qualname = getattr(func, "__qualname__", repr(func))
+
+    def _extract(hint: Any, where: str) -> T | None:
+        try:
+            return extract(hint)
+        except ValueError as exc:
+            exc.add_note(f"in {qualname} {where}")
+            raise
+
     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
+        if (payload := _extract(hint, f"input {name!r}")) is not None
     }
 
-    output: dict[str, _T] | _T | None
-    if is_typeddict(ret):
-        ret_hints = get_type_hints(ret, include_extras=True)
+    output: dict[str, T] | T | None = _extract(ret, "output")
+    record, _ = _record_return(ret)
+    if is_typeddict(record):
+        ret_hints = get_type_hints(record, include_extras=True)
         output = {
             name: payload
             for name, hint in ret_hints.items()
-            if (payload := extract(hint)) is not None
+            if (payload := _extract(hint, f"output {name!r}")) is not None
         }
-    elif dataclasses.is_dataclass(ret):
-        ret_hints = get_type_hints(ret, include_extras=True)
+    elif record is not None:
+        ret_hints = get_type_hints(record, include_extras=True)
         output = {
             f.name: payload
-            for f in dataclasses.fields(ret)
-            if (payload := extract(ret_hints.get(f.name))) is not None
+            for f in dataclasses.fields(record)
+            if (payload := _extract(ret_hints.get(f.name), f"output {f.name!r}"))
+            is not None
         }
-    else:
-        output = extract(ret)
 
     return inputs, output

xarray_annotated/_config.py

@@ -1,66 +1,33 @@
-"""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.
+"""The package-wide master switch, the only global that changes how checks behave.
+
+Behaviour is never configured globally, so that how a function validates can be read
+off its signature.  What remains global is whether xarray-annotated does anything at
+all.  The unit registry is process-wide too, but it decides which units parse, rather
+than what a check does with them.
+
+The switch resolves on every call as: an `enabled(...)` block, then the
+`XARRAY_ANNOTATED_ENABLED` environment variable, then on.  Code beats the
+environment in both directions, so `enabled(True)` restores checking in a process
+started with the variable set to false.
 """
 
 import os
-from collections.abc import Callable
-from typing import TypeVar
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+import numpy as np
 
-#: Master switch default. Validation is on unless explicitly disabled.
-DEFAULT_ENABLED: bool = True
+__all__ = ["ENABLED_ENV_VAR", "enabled", "resolve_enabled"]
 
-#: 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
+_override: ContextVar[bool | None] = ContextVar(
+    "xarray_annotated_enabled", default=None
+)
 
 
 def _parse_enabled_env(value: str) -> bool:
@@ -76,26 +43,58 @@ def _parse_enabled_env(value: str) -> bool:
 
 
 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
-    )
-
+    """Return whether xarray-annotated is active for this call.
 
-def get_enabled_override() -> bool | None:
-    """Return the current process override for the master switch (`None` if unset).
+    Returns:
+        The innermost `enabled(...)` value if inside one, else the parsed
+        `XARRAY_ANNOTATED_ENABLED`, else `True`.
 
-    Used by each domain's `policy` context manager to save/restore the shared
-    switch around a temporary override.
+    Raises:
+        ValueError: If the environment variable holds an unrecognised value.
     """
-    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.
+    override = _override.get()
+    if override is not None:
+        return override
+    env = os.environ.get(ENABLED_ENV_VAR)
+    if env:
+        return _parse_enabled_env(env)
+    return True
+
+
+@contextmanager
+def enabled(value: bool | np.bool_) -> Generator[None]:
+    """Switch xarray-annotated on or off for a block.
+
+    Off means no checks, so a body declaring `"Pa"` can receive `"hPa"`.  Declarations
+    are still validated when a function is decorated.
+
+    The block beats `XARRAY_ANNOTATED_ENABLED` in both directions, and is
+    scoped to the current thread or async task.  The variable accepts `1`,
+    `true`, `yes` or `on` and `0`, `false`, `no` or `off`, in any case; anything
+    else raises `ValueError` when a checked function is called.
+
+    Args:
+        value: `False` to switch off, `True` to switch on.
+
+    Raises:
+        TypeError: If `value` is not a bool, as `bool("false")` would leave checks on.
+
+    Examples:
+        >>> from typing import Annotated
+        >>> import xarray as xr
+        >>> import xarray_annotated as xra
+        >>> @xra.enforce
+        ... def f(p: Annotated[xr.DataArray, xra.Unit("Pa")]):
+        ...     return p
+        >>> hpa = xr.DataArray([1.0], attrs={"units": "hPa"})
+        >>> with xra.enabled(False):
+        ...     f(hpa).attrs["units"]
+        'hPa'
     """
-    global _process_enabled
-    _process_enabled = None if value is None else bool(value)
+    if not isinstance(value, bool | np.bool_):
+        raise TypeError(f"enabled takes a bool, not {type(value).__name__}.")
+    token = _override.set(bool(value))
+    try:
+        yield
+    finally:
+        _override.reset(token)