add a module-global registry (#32)

* use the application registry instead of creating a new one

* try to get the unit registry from units or existing quantities

* use a module-local registry instead of the application registry

* fix get_registry

* fix the registry property of DataArray

* remove the registry_kwargs test

* merge the units in get_registry

* don't use the name of the DataArray to set the units

* remove a comma

* mention that quantify will reject units from multiple registries

* document overriding the default registry

* rename the default registry to default_registry

* load pint into the docs namespace

* also load pint_xarray

* add a function set up a registry for use with pint_xarray

* replace the default registry with the application registry

* simplify the setup function

* update the docstring

* update whats-new.rst

* mention that the default registry is the application registry

Author: keewis

docs/creation.rst

@@ -8,6 +8,8 @@ Attaching units
 .. ipython:: python
     :suppress:
 
+    import pint
+    import pint_xarray
     import xarray as xr
 
 Usually, when loading data from disk we get a :py:class:`Dataset` or
@@ -40,9 +42,9 @@ We can also override the units of a variable:
 
     In [4]: ds.pint.quantify(b="km")
 
-    In [5]: da.pint.quantify({"b": "degree"})
+    In [5]: da.pint.quantify("degree")
 
-Overriding works, even if there is no ``units`` attribute, so we could use this
+Overriding works even if there is no ``units`` attribute, so we could use this
 to attach units to a normal :py:class:`Dataset`:
 
 .. ipython::
@@ -53,6 +55,13 @@ to attach units to a normal :py:class:`Dataset`:
 Of course, we could use :py:class:`pint.Unit` instances instead of strings to
 specify units, too.
 
+.. note::
+
+    Unit objects tied to different registries cannot interact with each
+    other. In order to avoid this, :py:meth:`DataArray.pint.quantify` and
+    :py:meth:`Dataset.pint.quantify` will make sure only a single registry is
+    used per ``xarray`` object.
+
 If we wanted to change the units of the data of a :py:class:`DataArray`, we
 could do so using the :py:attr:`DataArray.name` attribute:
 
@@ -76,6 +85,32 @@ have to first swap the dimensions:
        ...: )
        ...: da_with_units
 
+By default, :py:meth:`Dataset.pint.quantify` and
+:py:meth:`DataArray.pint.quantify` will use the unit registry at
+:py:obj:`pint_xarray.unit_registry` (the
+:py:func:`application registry <pint.get_application_registry>`). If we want a
+different registry, we can either pass it as the ``unit_registry`` parameter:
+
+.. ipython::
+
+   In [10]: ureg = pint.UnitRegistry(force_ndarray_like=True)
+       ...: # set up the registry
+
+   In [11]: da.pint.quantify("degree", unit_registry=ureg)
+
+or overwrite the default registry:
+
+.. ipython::
+
+   In [12]: pint_xarray.unit_registry = ureg
+
+   In [13]: da.pint.quantify("degree")
+
+.. note::
+
+    To properly work with ``xarray``, the ``force_ndarray_like`` or
+    ``force_ndarray`` options have to be enabled on the custom registry.
+
 Saving with units
 -----------------
 In order to not lose the units when saving to disk, we first have to call the

docs/whats-new.rst

@@ -13,3 +13,4 @@ What's new
 - fix the :py:attr:`DataArray.pint.units`, :py:attr:`DataArray.pint.magnitude`
   and :py:attr:`DataArray.pint.dimensionality` properties and add docstrings for
   all three. (:pull:`31`)
+- use ``pint``'s application registry as a module-global registry (:pull:`32`)

pint_xarray/__init__.py

@@ -8,6 +8,7 @@
 from . import testing  # noqa: F401
 from . import formatting
 from .accessors import PintDataArrayAccessor, PintDatasetAccessor  # noqa: F401
+from .accessors import default_registry as unit_registry  # noqa: F401
 
 try:
     __version__ = version("pint-xarray")

pint_xarray/accessors.py

@@ -8,6 +8,26 @@
 
 from . import conversion
 
+
+def setup_registry(registry):
+    """set up the given registry for use with pint_xarray
+
+    Namely, it enables ``force_ndarray_like`` to make sure results are always
+    duck arrays.
+
+    Parameters
+    ----------
+    registry : pint.UnitRegistry
+        The registry to modify
+    """
+    if not registry.force_ndarray and not registry.force_ndarray_like:
+        registry.force_ndarray_like = True
+
+    return registry
+
+
+default_registry = setup_registry(pint.get_application_registry())
+
 # TODO could/should we overwrite xr.open_dataset and xr.open_mfdataset to make
 # them apply units upon loading???
 # TODO could even override the decode_cf kwarg?
@@ -47,6 +67,16 @@ def zip_mappings(*mappings, fill_value=None):
     return zipped
 
 
+def merge_mappings(first, *mappings):
+    result = first.copy()
+    for mapping in mappings:
+        result.update(
+            {key: value for key, value in mapping.items() if value is not None}
+        )
+
+    return result
+
+
 def units_to_str_or_none(mapping):
     return {
         key: str(value) if isinstance(value, Unit) else value
@@ -72,13 +102,27 @@ def either_dict_or_kwargs(positional, keywords, method_name):
         return keywords
 
 
-def _get_registry(unit_registry, registry_kwargs):
+def get_registry(unit_registry, new_units, existing_units):
+    units = merge_mappings(existing_units, new_units)
+    registries = {unit._REGISTRY for unit in units.values() if isinstance(unit, Unit)}
+
     if unit_registry is None:
-        if registry_kwargs is None:
-            registry_kwargs = {}
-        registry_kwargs.update(force_ndarray=True)
-        # TODO should this registry object then be stored somewhere global?
-        unit_registry = pint.UnitRegistry(**registry_kwargs)
+        if not registries:
+            unit_registry = default_registry
+        elif len(registries) == 1:
+            (unit_registry,) = registries
+    registries.add(unit_registry)
+
+    if len(registries) > 1 or unit_registry not in registries:
+        raise ValueError(
+            "using multiple unit registries in the same object is not supported"
+        )
+
+    if not unit_registry.force_ndarray_like and not unit_registry.force_ndarray:
+        raise ValueError(
+            "invalid registry. Please enable 'force_ndarray_like' or 'force_ndarray'."
+        )
+
     return unit_registry
 
 
@@ -110,9 +154,7 @@ class PintDataArrayAccessor:
     def __init__(self, da):
         self.da = da
 
-    def quantify(
-        self, units=None, unit_registry=None, registry_kwargs=None, **unit_kwargs
-    ):
+    def quantify(self, units=None, unit_registry=None, **unit_kwargs):
         """
         Attaches units to the DataArray.
 
@@ -145,8 +187,6 @@ def quantify(
         unit_registry : pint.UnitRegistry, optional
             Unit registry to be used for the units attached to this DataArray.
             If not given then a default registry will be created.
-        registry_kwargs : dict, optional
-            Keyword arguments to be passed to the unit registry.
         **unit_kwargs
             Keyword argument form of units.
 
@@ -187,7 +227,11 @@ def quantify(
 
         units = either_dict_or_kwargs(units, unit_kwargs, ".quantify")
 
-        registry = _get_registry(unit_registry, registry_kwargs)
+        registry = get_registry(
+            unit_registry,
+            units,
+            conversion.extract_units(self.da),
+        )
 
         unit_attrs = conversion.extract_unit_attributes(self.da)
         new_obj = conversion.strip_unit_attributes(self.da)
@@ -254,7 +298,7 @@ def dimensionality(self):
     @property
     def registry(self):
         # TODO is this a bad idea? (see GH issue #1071 in pint)
-        return self.data._REGISTRY
+        return getattr(self.da.data, "_REGISTRY", None)
 
     @registry.setter
     def registry(self, _):
@@ -385,9 +429,7 @@ class PintDatasetAccessor:
     def __init__(self, ds):
         self.ds = ds
 
-    def quantify(
-        self, units=None, unit_registry=None, registry_kwargs=None, **unit_kwargs
-    ):
+    def quantify(self, units=None, unit_registry=None, **unit_kwargs):
         """
         Attaches units to each variable in the Dataset.
 
@@ -420,8 +462,6 @@ def quantify(
             Unit registry to be used for the units attached to each
             DataArray in this Dataset. If not given then a default
             registry will be created.
-        registry_kwargs : dict, optional
-            Keyword arguments to be passed to `pint.UnitRegistry`.
         **unit_kwargs
             Keyword argument form of ``units``.
 
@@ -458,7 +498,11 @@ def quantify(
             b        (x) int64 <Quantity([ 5 -2  1], 'decimeter')>
         """
         units = either_dict_or_kwargs(units, unit_kwargs, ".quantify")
-        registry = _get_registry(unit_registry, registry_kwargs)
+        registry = get_registry(
+            unit_registry,
+            units,
+            conversion.extract_units(self.ds),
+        )
 
         unit_attrs = conversion.extract_unit_attributes(self.ds)
         new_obj = conversion.strip_unit_attributes(self.ds)

pint_xarray/tests/test_accessors.py

@@ -94,11 +94,6 @@ def test_error_on_nonsense_units(self, example_unitless_da):
         with pytest.raises(UndefinedUnitError):
             da.pint.quantify(units="aecjhbav")
 
-    def test_registry_kwargs(self, example_unitless_da):
-        orig = example_unitless_da
-        result = orig.pint.quantify(registry_kwargs={"auto_reduce_dimensions": True})
-        assert result.data._REGISTRY.auto_reduce_dimensions is True
-
 
 class TestDequantifyDataArray:
     def test_strip_units(self, example_quantity_da):