xarray-annotated 0.6.0__tar.gz → 0.7.0__tar.gz

{xarray_annotated-0.6.0 → xarray_annotated-0.7.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: xarray-annotated
-Version: 0.6.0
-Summary: Run-time validation of xarray DataArray properties (units, dims/coords/dtype) declared via typing.Annotated
+Version: 0.7.0
+Summary: Run-time validation of xarray DataArray structure, units and frequency, declared via typing.Annotated
 Keywords: xarray,pint,units,validation,annotations
 Author: jmarshrossney
 Author-email: jmarshrossney <17361029+jmarshrossney@users.noreply.github.com>
@@ -33,21 +33,19 @@ Description-Content-Type: text/markdown
 [![Docs](https://github.com/jmarshrossney/xarray-annotated/actions/workflows/docs.yml/badge.svg)](https://jmarshrossney.github.io/xarray-annotated)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/jmarshrossney/xarray-annotated/blob/main/LICENSE)
 
-`xarray-annotated` enables run-time validation of `xarray.DataArray` properties declared in function signatures via [`typing.Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated).
+`xarray-annotated` checks `xarray.DataArray` arguments and return values against properties declared in the function signature with [`typing.Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated):
 
-The main idea is that annotations serve as a single source of truth for
+- **structure**: dimensions, coordinates and dtype;
+- **physical units**;
+- **time frequency** and phase.
 
-- structural schema (dims, coords, dtype),
-- physical units,
-- frequency (and phase) of a time axis,
+Making signatures the single source of truth has several advantages: the same declaration
 
-from which
+- can be used by automated checks and validation, and hence enforced at run time;
+- appears in generated documentation without needing to be hand-written;
+- provides rich context for AI coding agents.
 
-- static checkers and run-time validation can derive expected properties,
-- documentation can be generated automatically,
-- automated tools (coding agents) receive steering.
-
-For validating and automatically converting physical units we lean on the excellent [pint](https://pint.readthedocs.io/en/stable/) and [pint-xarray](https://pint-xarray.readthedocs.io).
+For physical units we lean on the excellent [pint](https://pint.readthedocs.io/en/stable/) and [pint-xarray](https://pint-xarray.readthedocs.io).
 ([cf-xarray](https://cf-xarray.readthedocs.io) is an optional dependency for CF/UDUNITS unit strings.)
 
 For full user documentation please visit **[https://jmarshrossney.github.io/xarray-annotated/](https://jmarshrossney.github.io/xarray-annotated/)**.
@@ -98,26 +96,37 @@ Attributes:
 The input units are wrong (hPa vs Pa) but the onus is on the user to check, and in reality we all know that such bugs can go unnoticed for a terrifyingly long time.
 The returned altitude is numerically wrong and carries a completely bogus hPa units in `attrs`.
 
-Here we use `xarray-annotated` to declare expected properties in the signature and enforce them at run-time with minimal boilerplate (just decorators).
+Here we use `xarray-annotated` to declare expected properties in the signature and enforce them at run-time with minimal boilerplate (one decorator).
 
+<!-- docs-test: raises ValueError -->
 ```python
 from typing import Annotated
 import xarray as xr
-from xarray_annotated.schema import declare_schema, Dims
-from xarray_annotated.units import declare_units, Unit
+import xarray_annotated as xra
+from xarray_annotated import Dims, Unit
 
 
-@declare_units
-@declare_schema
+@xra.enforce
 def altitude(
     p: Annotated[xr.DataArray, Dims("time"), Unit("Pa")],
 ) -> Annotated[xr.DataArray, Dims("time"), Unit("m")]:
     """Barometric altitude."""
-    return 44330.0 * (1.0 - (p / 101325.0) ** 0.1903)
+    alt = 44330.0 * (1.0 - (p / 101325.0) ** 0.1903)
+    return alt.assign_attrs(units="m")
 
 
 p = xr.DataArray([1013.0, 1000.0], dims=["time"], attrs={"units": "hPa"})
 altitude(p)
+# ValueError: [altitude] input 'p': unit 'hPa' differs from declared 'Pa'.
+#   Converting would change the values; if that is correct for this quantity,
+#   pass xra.units.convert(da, 'Pa').
+```
+
+The wrong unit is now an error rather than a wrong answer.
+Pressure is safe to convert, so we can do that explicitly:
+
+```python
+altitude(xra.units.convert(p, "Pa"))
 ```
 ```
 <xarray.DataArray (time: 2)> Size: 16B
@@ -127,34 +136,31 @@ Attributes:
     units:    m
 ```
 
-Now the outputs are physical and carry the correct units attribute.
-The `hPa` input was converted to `Pa` before the body ran, and the result was stamped with the unit the signature promised.
+Now the outputs are physical and carry the unit the signature promised.
+The body has to state that unit, because xarray would otherwise leave the `Pa` label on the result, and a label that disagrees with the declaration raises: it could just as easily mean a forgotten conversion.
+Nothing is [converted automatically](https://jmarshrossney.github.io/xarray-annotated/guides/units.html#converting-an-input), because whether a conversion is correct depends on the quantity, and a compatible unit is not always the right one.
 The docstring no longer needs to say "assumes Pa", because the signature does --- and unlike the docstring, it is actually enforced.
 
-If we had handed the function an array whose units were _dimensionally_ wrong, or with an extra dimension, this would have failed gracefully:
+An array whose units were _dimensionally_ wrong would also have raised:
 
+<!-- docs-test: raises DimensionalityError -->
 ```python
 altitude(xr.DataArray([1.0, 2.0], dims=["time"], attrs={"units": "kg"}))
 # DimensionalityError: Cannot convert from 'kilogram' ([mass])
 #   to 'pascal' ([mass] / [length] / [time] ** 2)
-
-altitude(xr.DataArray([[1013.0, 1000.0]], dims=["run", "time"], attrs={"units": "hPa"}))
-# SchemaError: [altitude] 'p' dims mismatch: expected ('time',) in any order, got ('run', 'time')
 ```
 
-## Contributing
+And so would one with an extra dimension:
 
-```sh
-uv sync   # install the package + dev dependencies into .venv
-just      # lint, typecheck, test, build docs
+<!-- docs-test: raises SchemaError -->
+```python
+altitude(
+    xr.DataArray([[101300.0, 100000.0]], dims=["run", "time"], attrs={"units": "Pa"})
+)
+# SchemaError: [altitude] 'p' dims mismatch: expected ('time',) in any order, got ('run', 'time')
 ```
 
-See the `justfile` for individual targets (`just lint`, `just test`, `just test-cov`, ...).
-
-A `.pre-commit-config.yaml` is included to run the same linting (ruff) and type-checking (pyright) steps on every commit. Install the hooks with:
-
-```sh
-uv run pre-commit install
-```
+## Contributing
 
-Contributions are welcome --- open an issue or pull request on GitHub.
+Issues and pull requests are welcome.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for the package's scope and how to set up a development environment.

{xarray_annotated-0.6.0 → xarray_annotated-0.7.0}/README.md

@@ -6,21 +6,19 @@
 [![Docs](https://github.com/jmarshrossney/xarray-annotated/actions/workflows/docs.yml/badge.svg)](https://jmarshrossney.github.io/xarray-annotated)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/jmarshrossney/xarray-annotated/blob/main/LICENSE)
 
-`xarray-annotated` enables run-time validation of `xarray.DataArray` properties declared in function signatures via [`typing.Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated).
+`xarray-annotated` checks `xarray.DataArray` arguments and return values against properties declared in the function signature with [`typing.Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated):
 
-The main idea is that annotations serve as a single source of truth for
+- **structure**: dimensions, coordinates and dtype;
+- **physical units**;
+- **time frequency** and phase.
 
-- structural schema (dims, coords, dtype),
-- physical units,
-- frequency (and phase) of a time axis,
+Making signatures the single source of truth has several advantages: the same declaration
 
-from which
+- can be used by automated checks and validation, and hence enforced at run time;
+- appears in generated documentation without needing to be hand-written;
+- provides rich context for AI coding agents.
 
-- static checkers and run-time validation can derive expected properties,
-- documentation can be generated automatically,
-- automated tools (coding agents) receive steering.
-
-For validating and automatically converting physical units we lean on the excellent [pint](https://pint.readthedocs.io/en/stable/) and [pint-xarray](https://pint-xarray.readthedocs.io).
+For physical units we lean on the excellent [pint](https://pint.readthedocs.io/en/stable/) and [pint-xarray](https://pint-xarray.readthedocs.io).
 ([cf-xarray](https://cf-xarray.readthedocs.io) is an optional dependency for CF/UDUNITS unit strings.)
 
 For full user documentation please visit **[https://jmarshrossney.github.io/xarray-annotated/](https://jmarshrossney.github.io/xarray-annotated/)**.
@@ -71,26 +69,37 @@ Attributes:
 The input units are wrong (hPa vs Pa) but the onus is on the user to check, and in reality we all know that such bugs can go unnoticed for a terrifyingly long time.
 The returned altitude is numerically wrong and carries a completely bogus hPa units in `attrs`.
 
-Here we use `xarray-annotated` to declare expected properties in the signature and enforce them at run-time with minimal boilerplate (just decorators).
+Here we use `xarray-annotated` to declare expected properties in the signature and enforce them at run-time with minimal boilerplate (one decorator).
 
+<!-- docs-test: raises ValueError -->
 ```python
 from typing import Annotated
 import xarray as xr
-from xarray_annotated.schema import declare_schema, Dims
-from xarray_annotated.units import declare_units, Unit
+import xarray_annotated as xra
+from xarray_annotated import Dims, Unit
 
 
-@declare_units
-@declare_schema
+@xra.enforce
 def altitude(
     p: Annotated[xr.DataArray, Dims("time"), Unit("Pa")],
 ) -> Annotated[xr.DataArray, Dims("time"), Unit("m")]:
     """Barometric altitude."""
-    return 44330.0 * (1.0 - (p / 101325.0) ** 0.1903)
+    alt = 44330.0 * (1.0 - (p / 101325.0) ** 0.1903)
+    return alt.assign_attrs(units="m")
 
 
 p = xr.DataArray([1013.0, 1000.0], dims=["time"], attrs={"units": "hPa"})
 altitude(p)
+# ValueError: [altitude] input 'p': unit 'hPa' differs from declared 'Pa'.
+#   Converting would change the values; if that is correct for this quantity,
+#   pass xra.units.convert(da, 'Pa').
+```
+
+The wrong unit is now an error rather than a wrong answer.
+Pressure is safe to convert, so we can do that explicitly:
+
+```python
+altitude(xra.units.convert(p, "Pa"))
 ```
 ```
 <xarray.DataArray (time: 2)> Size: 16B
@@ -100,34 +109,31 @@ Attributes:
     units:    m
 ```
 
-Now the outputs are physical and carry the correct units attribute.
-The `hPa` input was converted to `Pa` before the body ran, and the result was stamped with the unit the signature promised.
+Now the outputs are physical and carry the unit the signature promised.
+The body has to state that unit, because xarray would otherwise leave the `Pa` label on the result, and a label that disagrees with the declaration raises: it could just as easily mean a forgotten conversion.
+Nothing is [converted automatically](https://jmarshrossney.github.io/xarray-annotated/guides/units.html#converting-an-input), because whether a conversion is correct depends on the quantity, and a compatible unit is not always the right one.
 The docstring no longer needs to say "assumes Pa", because the signature does --- and unlike the docstring, it is actually enforced.
 
-If we had handed the function an array whose units were _dimensionally_ wrong, or with an extra dimension, this would have failed gracefully:
+An array whose units were _dimensionally_ wrong would also have raised:
 
+<!-- docs-test: raises DimensionalityError -->
 ```python
 altitude(xr.DataArray([1.0, 2.0], dims=["time"], attrs={"units": "kg"}))
 # DimensionalityError: Cannot convert from 'kilogram' ([mass])
 #   to 'pascal' ([mass] / [length] / [time] ** 2)
-
-altitude(xr.DataArray([[1013.0, 1000.0]], dims=["run", "time"], attrs={"units": "hPa"}))
-# SchemaError: [altitude] 'p' dims mismatch: expected ('time',) in any order, got ('run', 'time')
 ```
 
-## Contributing
+And so would one with an extra dimension:
 
-```sh
-uv sync   # install the package + dev dependencies into .venv
-just      # lint, typecheck, test, build docs
+<!-- docs-test: raises SchemaError -->
+```python
+altitude(
+    xr.DataArray([[101300.0, 100000.0]], dims=["run", "time"], attrs={"units": "Pa"})
+)
+# SchemaError: [altitude] 'p' dims mismatch: expected ('time',) in any order, got ('run', 'time')
 ```
 
-See the `justfile` for individual targets (`just lint`, `just test`, `just test-cov`, ...).
-
-A `.pre-commit-config.yaml` is included to run the same linting (ruff) and type-checking (pyright) steps on every commit. Install the hooks with:
-
-```sh
-uv run pre-commit install
-```
+## Contributing
 
-Contributions are welcome --- open an issue or pull request on GitHub.
+Issues and pull requests are welcome.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for the package's scope and how to set up a development environment.

{xarray_annotated-0.6.0 → xarray_annotated-0.7.0}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "uv_build"
 
 [project]
 name = "xarray-annotated"
-version = "0.6.0"
-description = "Run-time validation of xarray DataArray properties (units, dims/coords/dtype) declared via typing.Annotated"
+version = "0.7.0"
+description = "Run-time validation of xarray DataArray structure, units and frequency, declared via typing.Annotated"
 readme = "README.md"
 license = "MIT"
 license-files = ["LICENSE"]
@@ -60,7 +60,9 @@ dev = [
 ]
 integrations = [
     "beartype>=0.22.9",
+    "dask>=2026.8.0",
     "jaxtyping>=0.3.11",
+    "netcdf4>=1.7.5.1.post1",
 ]
 docs = [
     "marimo>=0.23.11",

{xarray_annotated-0.6.0 → xarray_annotated-0.7.0}/pyproject.toml.orig

@@ -4,8 +4,8 @@ build-backend = "uv_build"
 
 [project]
 name = "xarray-annotated"
-version = "0.6.0"
-description = "Run-time validation of xarray DataArray properties (units, dims/coords/dtype) declared via typing.Annotated"
+version = "0.7.0"
+description = "Run-time validation of xarray DataArray structure, units and frequency, declared via typing.Annotated"
 authors = [
     { name = "jmarshrossney", email = "17361029+jmarshrossney@users.noreply.github.com" }
 ]
@@ -55,10 +55,13 @@ dev = [
 # test-only: there is nothing for a user to install to "enable" the integration, so this
 # is deliberately a dependency group and not an optional-dependency extra (contrast
 # `cf`, which is an extra because `units._registry` really does import cf_xarray).
-# `tests/test_jaxtyping.py` skips itself when they are absent.
+# `tests/test_jaxtyping.py`, `test_dask.py` and `test_netcdf.py` skip themselves when
+# they are absent.
 integrations = [
     "beartype>=0.22.9",
+    "dask>=2026.8.0",
     "jaxtyping>=0.3.11",
+    "netcdf4>=1.7.5.1.post1",
 ]
 docs = [
     "marimo>=0.23.11",

xarray_annotated-0.7.0/src/xarray_annotated/__init__.py

@@ -0,0 +1,60 @@
+"""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 enforced at runtime.  Each kind of
+property is a domain subpackage, imported explicitly:
+
+* ``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.
+
+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 ._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",
+]

xarray_annotated-0.7.0/src/xarray_annotated/_annotations.py

@@ -0,0 +1,179 @@
+"""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. 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, Iterable
+from typing import (
+    Annotated,
+    Any,
+    TypeAliasType,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+    is_typeddict,
+)
+
+import xarray as xr
+
+
+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.
+
+    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.
+
+    Raises:
+        ValueError: If `metadata` holds more than one `marker_type` marker.
+    """
+    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 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, 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 _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 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 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 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, 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, f"input {name!r}")) is not None
+    }
+
+    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, f"output {name!r}")) is not None
+        }
+    elif record is not None:
+        ret_hints = get_type_hints(record, include_extras=True)
+        output = {
+            f.name: payload
+            for f in dataclasses.fields(record)
+            if (payload := _extract(ret_hints.get(f.name), f"output {f.name!r}"))
+            is not None
+        }
+
+    return inputs, output

xarray_annotated-0.7.0/src/xarray_annotated/_config.py

@@ -0,0 +1,100 @@
+"""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 Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+import numpy as np
+
+__all__ = ["ENABLED_ENV_VAR", "enabled", "resolve_enabled"]
+
+ENABLED_ENV_VAR = "XARRAY_ANNOTATED_ENABLED"
+
+_TRUTHY: frozenset[str] = frozenset({"1", "true", "yes", "on"})
+_FALSEY: frozenset[str] = frozenset({"0", "false", "no", "off"})
+
+_override: ContextVar[bool | None] = ContextVar(
+    "xarray_annotated_enabled", default=None
+)
+
+
+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:
+    """Return whether xarray-annotated is active for this call.
+
+    Returns:
+        The innermost `enabled(...)` value if inside one, else the parsed
+        `XARRAY_ANNOTATED_ENABLED`, else `True`.
+
+    Raises:
+        ValueError: If the environment variable holds an unrecognised value.
+    """
+    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'
+    """
+    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)