Get pint unit registries from the input Quantities

Author: Andrew Barna

gsw_xarray/_core.py

@@ -12,8 +12,6 @@
     import pint_xarray
     import pint
 
-    ureg = pint.UnitRegistry()
-    Q_ = ureg.Quantity
 except ImportError:
     pint_xarray = None
 
@@ -24,75 +22,84 @@ def add_attrs(rv, attrs, name):
         rv.attrs = attrs
 
 
-def quantify(rv, attrs):
+def quantify(rv, attrs, unit_registry=None):
+    if unit_registry is None:
+        return rv
+
     if isinstance(rv, xr.DataArray):
-        rv = rv.pint.quantify()
+        rv = rv.pint.quantify(unit_registry=unit_registry)
     else:
         if attrs is not None:
             # Necessary to use the Q_ and not simply multiplication with ureg unit because of temperature
             # see https://pint.readthedocs.io/en/latest/nonmult.html
-            rv = Q_(rv, attrs["units"])
+            rv = unit_registry.Quantity(rv, attrs["units"])
     return rv
 
 
 def pint_compat(args, kwargs):
     if pint_xarray is None:
-        return args, kwargs, False
+        return args, kwargs, None
 
     using_pint = False
     new_args = []
     new_kwargs = {}
+    registries = []
     for arg in args:
         if isinstance(arg, xr.DataArray):
             if arg.pint.units is not None:
                 new_args.append(arg.pint.dequantify())
-                using_pint = True
+                registries.append(arg.pint.registry)
             else:
                 new_args.append(arg)
         elif isinstance(arg, pint.Quantity):
             new_args.append(arg.magnitude)
-            using_pint = True
+            registries.append(arg._REGISTRY)
         else:
             new_args.append(arg)
 
     for kw, arg in kwargs.items():
         if isinstance(arg, xr.DataArray):
             if arg.pint.units is not None:
                 new_kwargs[kw] = arg.pint.dequantify()
-                using_pint = True
+                registries.append(arg.pint.registry)
             else:
                 new_kwargs[kw] = arg
         elif isinstance(arg, pint.Quantity):
             new_kwargs[kw] = arg.magnitude
-            using_pint = True
+            registries.append(arg._REGISTRY)
         else:
             new_kwargs[kw] = arg
 
-    return new_args, new_kwargs, using_pint
+    registries = set(registries)
+    if len(registries) > 1:
+        raise ValueError("Quantity arguments must all belong to the same unit registry")
+    elif len(registries) == 0:
+        registries = None
+    else:
+        (registries,) = registries
+    return new_args, new_kwargs, registries
 
 
 def cf_attrs(attrs, name, check_func):
     def cf_attrs_decorator(func):
         @wraps(func)
         def cf_attrs_wrapper(*args, **kwargs):
-            args, kwargs, is_quantity = pint_compat(args, kwargs)
+            args, kwargs, unit_registry = pint_compat(args, kwargs)
             rv = func(*args, **kwargs)
             attrs_checked = check_func(attrs, args, kwargs)
             if isinstance(rv, tuple):
                 rv_updated = []
                 for (i, da) in enumerate(rv):
                     add_attrs(da, attrs_checked[i], name[i])
-                    if is_quantity:
-                        rv_updated.append(quantify(da, attrs_checked[i]))
-                    else:
-                        rv_updated.append(da)
+                    rv_updated.append(
+                        quantify(da, attrs_checked[i], unit_registry=unit_registry)
+                    )
 
                 rv = tuple(rv_updated)
 
             else:
                 add_attrs(rv, attrs_checked, name)
-                if is_quantity:
-                    rv = quantify(rv, attrs_checked)
+                rv = quantify(rv, attrs_checked, unit_registry=unit_registry)
             return rv
 
         return cf_attrs_wrapper

gsw_xarray/tests/test_units.py

@@ -85,9 +85,8 @@ def test_func_return_tuple_quantity(ds_pint):
 def test_pint_quantity_xarray(ds):
     """If input is mixed between xr.DataArray and pint quantity it should return pint-xarray wrapped quantity"""
     pint_xarray = pytest.importorskip("pint_xarray")
-    import pint
 
-    ureg = pint.UnitRegistry()
+    ureg = pint_xarray.unit_registry
     Q_ = ureg.Quantity
     sigma0 = gsw_xarray.sigma0(SA=ds.SA, CT=Q_(25.4, ureg.degC))
     assert sigma0.pint.units == pint_xarray.unit_registry("kg / m^3")
@@ -98,7 +97,7 @@ def test_pint_quantity():
     pint_xarray = pytest.importorskip("pint_xarray")
     import pint
 
-    ureg = pint.UnitRegistry()
+    ureg = pint_xarray.unit_registry
     CT = gsw_xarray.CT_from_pt(SA=35 * ureg("g / kg"), pt=10)
     assert isinstance(CT, pint.Quantity)
 
@@ -108,7 +107,20 @@ def test_pint_quantity_tuple():
     pint_xarray = pytest.importorskip("pint_xarray")
     import pint
 
-    ureg = pint.UnitRegistry()
+    ureg = pint_xarray.unit_registry
     (a, b) = gsw_xarray.CT_first_derivatives(35 * ureg("g / kg"), pt=1)
     assert isinstance(a, pint.Quantity)
     assert isinstance(b, pint.Quantity)
+
+
+def test_mixed_unit_regestiries():
+    """If input quantities are from different registries, it should fail"""
+    pint_xarray = pytest.importorskip("pint_xarray")
+    import pint
+
+    ureg_a = pint.UnitRegistry()
+    ureg_b = pint.UnitRegistry()
+    with pytest.raises(ValueError):
+        gsw_xarray.CT_first_derivatives(
+            35 * ureg_a("g / kg"), pt=ureg_b.Quantity(1, ureg_b.degC)
+        )