xarray_annotated.units¶

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.

Parameters:

• hint (Any) –

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

Returns:

• str | None –

  The declared unit (first Unit marker, else first str) if the base

• str | None –

  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

Apply a declared unit to a returned DataArray, per the on-output axis.

The output counterpart to check_units, and deliberately not symmetric with it: a declared output carrying an attrs label is stamped, never converted.

The reason is that xarray’s attrs are inert under arithmetic. A body that converts by scalar multiplication — return p * 100.0, or flux * F for a molar-mass factor — produces an array still labelled with its input’s unit, because multiplying by a float cannot update a string. So for the dominant idiom the returned label describes neither the values nor the declaration, and converting on the strength of it would rescale correct values a second time.

Two modes follow from that:

• "stamp" (default) — overwrite attrs["units"] with declared, no checks. The only sound default, because a stale label is the norm rather than a symptom.
• "strict" — a label that is present, parseable, and not equal to declared raises.

"strict" is only meaningful for bodies that maintain their own units — pass-through and subsetting functions, or computations on pint-quantified arrays. A body doing manual scalar arithmetic will fail it whether or not it is correct, so this is an opt-in for unit-aware code, not a general verification mode. (If a body must do manual arithmetic and you still want "strict" elsewhere, have it drop the stale label — attrs cleared, or keep_attrs=False — since an absent label is always stamped.)

Pint-quantified returns are the exception, and are converted rather than stamped. A Quantity’s unit lives in the data, so — unlike attrs — it cannot be left behind by arithmetic; it is always a true description of the values. Converting on the strength of it is therefore sound, and stamping on the strength of it would be actively wrong: writing attrs["units"] onto a Quantity that disagrees produces an array labelled two ways at once. So for a quantified return:

• "stamp" — converts to declared if it is not already there, and raises pint.DimensionalityError if it cannot. This is the one case where "stamp" can raise, because here a mismatch is a real error rather than the expected consequence of an inert label.
• "strict" — a unit other than declared raises, as for attrs.

Either way the array is returned still quantified, with attrs untouched: it already describes itself, and dequantifying it purely to write a string would copy the whole buffer.

Note what no mode can do: distinguish flux * F from flux when both are labelled the same. Output values are never inspected, so a forgotten conversion factor is not detectable here — only at a consumer that declares the quantity and receives an array that was never stamped.

An absent, over-long, or unparseable label is always stamped: absence is not evidence of a mismatch.

Parameters:

• da (DataArray) –

  The returned DataArray.

• declared (str) –

  The declared unit string (e.g. "Pa").

• name (str) –

  A label for the value, used in error messages.

• on_output (OnOutput | None, default: None ) –

  Override the on-output axis for this call (None defers to the active policy).

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

  A qualifier prepended to messages as [qualname].

Returns:

• DataArray –

  For an attrs-labelled or unlabelled array, da itself with

• DataArray –

  attrs["units"] set to declared (mutated in place). For a quantified

• DataArray –

  array, an array in declared with attrs untouched — the same object

• DataArray –

  when it was already in declared, otherwise a converted new one. Always

• DataArray –

  use the return value rather than relying on mutation.

Raises:

• DimensionalityError –

  Under "strict", when the returned array’s own unit is dimensionally incompatible with declared; also under "stamp" for a quantified array that cannot be converted.

• ValueError –

  Under "strict", when the returned array’s unit is compatible with but not equal to declared.

Examples:

>>> import xarray as xr
>>> from xarray_annotated.units._check import apply_output_units
>>> da = xr.DataArray([1.0])
>>> apply_output_units(da, "Pa", "return").attrs["units"]
'Pa'

Raise ValueError if unit is not parseable by the active registry.

Used to fail fast at declaration time: a malformed or undefined unit string (a typo such as "degrees_C", or "not_a_unit") is rejected as soon as it is declared, rather than only when it is later used to validate data.

Accepts either a bare unit string or a Unit marker (its .unit string is validated), so a consumer that read a Unit off an annotation can pass it straight through.

The registry raises a variety of exception types for bad input (pint.UndefinedUnitError, AssertionError, …); all are caught and re-raised as a single, clear ValueError.

Parameters:

• unit (str | Unit) –

  The unit string (or Unit marker) to validate.

• context (str) –

  Names the offending site (e.g. "mymodel input 'vpd_weekly'") for the error message.

Raises:

• ValueError –

  If unit cannot be parsed by the active registry.

Examples:

>>> from xarray_annotated.units import assert_valid_unit
>>> assert_valid_unit("Pa", "test")  # no raise
>>> assert_valid_unit("degC", "test")  # no raise

Validate and convert an input DataArray to its declared unit.

Returns a DataArray whose data is expressed in declared and whose units attribute equals declared.

Behaviour follows the active Policy (get_policy); on_missing and on_inexact override their axes for this call when given (None defers to the policy). If the policy is disabled the array is returned unchanged.

Three events are handled, each building on the one before:

• No parseable unit — the input has no units attribute, or one the registry cannot parse (e.g. a non-CF string like "fraction"): follows on_missing ("error" raises, "warn" warns and returns unchanged, "ignore" returns unchanged silently).
• Value-changing conversion — the unit is dimensionally compatible with declared but not the same unit (e.g. "hPa" where "Pa" is declared): follows on_inexact ("error" raises, "warn" warns then converts, "convert" converts silently). Equivalent spellings ("pascal" for "Pa") imply no value change and are simply relabelled, without touching the data.
• Dimensional mismatch — two parseable but incompatible units (e.g. a mass where a pressure is declared): always raises pint.DimensionalityError, regardless of policy.

Parameters:

• da (DataArray) –

  The input DataArray to validate and convert.

• declared (str) –

  The expected unit string (e.g. "Pa").

• name (str) –

  A label for the array, used in error/warning messages.

• on_missing (OnMissing | None, default: None ) –

  Override the on-missing axis for this call (None defers to the active policy).

• on_inexact (OnInexact | None, default: None ) –

  Override the on-inexact axis for this call (None defers to the active policy).

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

  A qualifier prepended to warning messages as [qualname], identifying the calling site.

Returns:

• DataArray –

  A new DataArray converted to declared, with attrs["units"] set

• DataArray –

  to declared. When the input is already in declared the data is not

• DataArray –

  copied — the result is a shallow copy sharing the caller’s buffer, so

• DataArray –

  declaring a unit an array already has costs nothing.

Raises:

• ValueError –

  When on_missing="error" and no parseable unit is found, or when on_inexact="error" and a value-changing conversion would occur.

• DimensionalityError –

  When the actual unit is dimensionally incompatible with declared (always, regardless of policy).

Examples:

>>> import numpy as np, xarray as xr
>>> from xarray_annotated.units import check_units
>>> da = xr.DataArray([1013.0, 1000.0], attrs={"units": "hPa"})
>>> out = check_units(da, "Pa", "pressure")
>>> out.attrs["units"]
'Pa'
>>> out.values  # 10 * 100 = 1000
array([101300., 100000.])

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

Reads the decorated function’s own type annotations once, via units_from_signature: parameters annotated Annotated[DataArray, Unit(...)] (or the bare-string shorthand) declare input units, and a TypedDict return (or a bare Annotated[DataArray, Unit(...)] return) declares output units. Those annotations are the single source of truth, so a unit is never written twice.

On each call, under the active Policy (get_policy), the wrapper:

1. validates/converts every declared DataArray input to its unit via check_units;
2. runs the wrapped function;
3. applies each declared output unit to the returned DataArray (a dict return is handled per key; a single DataArray return takes the bare declared unit).

Note the asymmetry in step 3: an output is stamped, not converted. A body that does its own unit arithmetic (return p * 100.0) leaves attrs holding the input’s unit, so converting on the strength of that label would scale correct values a second time. The on_output axis can require the label to match instead ("strict"), which suits only bodies that maintain their own units; see apply_output_units for the trade-off.

Only DataArray values are touched; other arguments and returns pass through unchanged. When the policy is disabled (enabled=False) the wrapper is a total no-op: inputs are not converted and outputs are not stamped.

Usable bare (@declare_units) or called (@declare_units(on_missing="error")).

Every declared unit string is checked against the registry at decoration time, so a malformed or undefined unit fails fast at import — regardless of policy — rather than only when the function first runs.

Parameters:

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

  The function to decorate (when used bare). None when called with keyword arguments, in which case a parametrised decorator is returned.

• on_missing (OnMissing | None, default: None ) –

  Override the on-missing axis for the decorated function. When None (default) resolved per call from get_policy.

• on_inexact (OnInexact | None, default: None ) –

  Override the on-inexact axis for the decorated function. When None (default) resolved per call from get_policy.

• on_output (OnOutput | None, default: None ) –

  Override the on-output axis for the decorated function. When None (default) resolved per call from get_policy.

Returns:

• Callable[..., Any] –

  A wrapped function that validates inputs and stamps outputs according

• Callable[..., Any] –

  to the active policy.

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

Returns:

• Policy –

  The resolved Policy.

Return the active pint UnitRegistry.

Defaults to pint.get_application_registry() (plain pint) until set_registry or use_cf_units is called.

Returns:

• UnitRegistry | ApplicationRegistry –

  The process-global pint UnitRegistry.

Temporarily override policy axes, restoring all of 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_missing (OnMissing | _Unset | None, default: _UNSET ) –

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

• on_inexact (OnInexact | _Unset | None, default: _UNSET ) –

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

• on_output (OnOutput | _Unset | None, default: _UNSET ) –

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

Examples:

>>> from xarray_annotated.units import policy
>>> import xarray as xr
>>> da = xr.DataArray([1.0], attrs={"units": "Pa"})
>>> with policy(on_missing="error"):
...     pass  # policy restored after block

Set process-wide 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_missing (OnMissing | _Unset | None, default: _UNSET ) –

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

• on_inexact (OnInexact | _Unset | None, default: _UNSET ) –

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

• on_output (OnOutput | _Unset | None, default: _UNSET ) –

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

Set the process-wide pint registry used by this module and pint-xarray.

Also calls pint_xarray.setup_registry(ureg) so the .pint accessor and this module’s parse/compat helpers never drift apart, and pint.set_application_registry(ureg) so pint’s global application registry (which pint-xarray follows for .pint.quantify()) stays in sync with this module’s choice. Without the latter, switching back to plain pint after use_cf_units would leave the application registry pointing at the CF one.

pint has a single process-global application registry, so this is a one-time, startup choice, not a per-array setting: quantities created under two different registries cannot be mixed (pint raises).

Parameters:

• ureg (UnitRegistry | ApplicationRegistry) –

  The UnitRegistry to install as the process-global registry.

Return whether two units are dimensionally compatible.

Mirrors the runtime conversion semantics of check_units: "hPa" and "Pa" are compatible (one converts to the other), whereas "Pa" and "kg" are not. Both are assumed already validated by assert_valid_unit.

Accepts either a bare unit string or a Unit marker on each side (its .unit string is used), so a consumer holding markers — e.g. Declared values from declarations_from_signature — can compare them directly.

Parameters:

• a (str | Unit) –

  A unit string or Unit marker.

• b (str | Unit) –

  Another unit string or Unit marker.

Returns:

• bool –

  True if the units are dimensionally compatible.

Examples:

>>> from xarray_annotated.units import Unit, units_compatible
>>> units_compatible("hPa", "Pa")
True
>>> units_compatible(Unit("hPa"), "Pa")  # markers accepted too
True
>>> units_compatible("Pa", "kg")
False

Return whether two units are the same unit (no conversion needed).

Compares the parsed units, so different spellings of the same unit are equal ("Pa" == "pascal", "1" == "dimensionless") while a prefixed unit differs ("hPa" != "Pa"). This is the distinction the on_inexact axis turns on: a value-changing conversion is one where the units are compatible but not equal; equivalent spellings imply no value change.

Accepts either a bare unit string or a Unit marker on each side (its .unit string is used), matching units_compatible.

Parameters:

• a (str | Unit) –

  A unit string or Unit marker.

• b (str | Unit) –

  Another unit string or Unit marker.

Returns:

• bool –

  True if the units are the same (no conversion needed).

Examples:

>>> from xarray_annotated.units import Unit, units_equal
>>> units_equal("Pa", "pascal")
True
>>> units_equal(Unit("Pa"), Unit("pascal"))  # markers accepted too
True
>>> units_equal("hPa", "Pa")
False

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.

Parameters:

• func (object) –

  A callable whose type hints carry Annotated[DataArray, Unit(...)] metadata (or the bare-string shorthand).

Returns:

• dict[str, str] –

  A (input_units, output_units) pair:

• dict[str, str] | str | None –
  • input_units — dict[str, str] mapping parameter names to their declared unit strings. Only parameters whose hint carries a unit declaration contribute.
• tuple[dict[str, str], dict[str, str] | str | None] –
  • 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.

Activate cf-xarray’s UDUNITS-aware registry.

Lazily imports cf_xarray.units (from the [cf] extra) and installs its registry via set_registry, so CF-convention unit strings such as "umol m-2 s-1" and "g m-2 d-1" parse.

Raises:

• ImportError –

  If cf-xarray is not installed.