add support for pint only quantities

Author: rcaneill

gsw_xarray/_core.py

@@ -11,6 +11,9 @@
 try:
     import pint_xarray
     import pint
+
+    ureg = pint.UnitRegistry()
+    Q_ = ureg.Quantity
 except ImportError:
     pint_xarray = None
 
@@ -20,6 +23,16 @@ def add_attrs(rv, attrs, name):
         rv.name = name
         rv.attrs = attrs
 
+
+def quantify(rv, attrs):
+    if isinstance(rv, xr.DataArray):
+        rv = rv.pint.quantify()
+    else:
+        if attrs is not None:
+            rv = Q_(rv, attrs["units"])
+    return rv
+
+
 def pint_compat(args, kwargs):
     if pint_xarray is None:
         return args, kwargs, False
@@ -68,16 +81,16 @@ def cf_attrs_wrapper(*args, **kwargs):
                 for (i, da) in enumerate(rv):
                     add_attrs(da, attrs_checked[i], name[i])
                     if is_quantity:
-                        rv_updated.append(da.pint.quantify())
+                        rv_updated.append(quantify(da, attrs_checked[i]))
                     else:
                         rv_updated.append(da)
 
                 rv = tuple(rv_updated)
-                
+
             else:
                 add_attrs(rv, attrs_checked, name)
                 if is_quantity:
-                    rv = rv.pint.quantify()
+                    rv = quantify(rv, attrs_checked)
             return rv
 
         return cf_attrs_wrapper

gsw_xarray/tests/test_units.py

@@ -50,30 +50,65 @@ def test_xarray_quantity(ds_pint):
     sigma0 = gsw_xarray.sigma0(SA=ds_pint.SA, CT=ds_pint.CT)
     assert sigma0.pint.units == pint_xarray.unit_registry("kg / m^3")
 
-@pytest.mark.parametrize("SA_type", ['unit', 'ds'])
-@pytest.mark.parametrize("CT_type", ['unit', 'ds'])
+
+@pytest.mark.parametrize("SA_type", ["unit", "ds"])
+@pytest.mark.parametrize("CT_type", ["unit", "ds"])
 def test_xarray_quantity_or_ds(ds, ds_pint, SA_type, CT_type):
     """If at least 1 of the inputs is quantity, the result should be quantity"""
     pint_xarray = pytest.importorskip("pint_xarray")
-    if SA_type == 'unit':
+    if SA_type == "unit":
         SA = ds_pint.SA
-    elif SA_type == 'ds':
+    elif SA_type == "ds":
         SA = ds.SA
-    
-    if CT_type == 'unit':
+
+    if CT_type == "unit":
         CT = ds_pint.CT
-    elif CT_type == 'ds':
+    elif CT_type == "ds":
         CT = ds.CT
-        
+
     sigma0 = gsw_xarray.sigma0(SA=SA, CT=CT)
-    if SA_type == 'unit' or CT_type == 'unit':
+    if SA_type == "unit" or CT_type == "unit":
         assert sigma0.pint.units == pint_xarray.unit_registry("kg / m^3")
     else:
         assert sigma0.pint.units is None
-        assert sigma0.pint.quantify().pint.units == pint_xarray.unit_registry("kg / m^3")
+        assert sigma0.pint.quantify().pint.units == pint_xarray.unit_registry(
+            "kg / m^3"
+        )
 
 
 def test_func_return_tuple_quantity(ds_pint):
     pint_xarray = pytest.importorskip("pint_xarray")
     (CT_SA, CT_pt) = gsw_xarray.CT_first_derivatives(ds_pint.SA, 1)
     assert CT_SA.pint.units == pint_xarray.unit_registry("K/(g/kg)")
+
+
+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()
+    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")
+
+
+def test_pint_quantity():
+    """If input is pint quantity should return a quantity"""
+    pint_xarray = pytest.importorskip("pint_xarray")
+    import pint
+
+    ureg = pint.UnitRegistry()
+    CT = gsw_xarray.CT_from_pt(SA=35 * ureg("g / kg"), pt=10)
+    assert isinstance(CT, pint.Quantity)
+
+
+def test_pint_quantity_tuple():
+    """If input is pint quantity should return a quantity"""
+    pint_xarray = pytest.importorskip("pint_xarray")
+    import pint
+
+    ureg = pint.UnitRegistry()
+    (a, b) = gsw_xarray.CT_first_derivatives(35 * ureg("g / kg"), pt=1)
+    assert isinstance(a, pint.Quantity)
+    assert isinstance(b, pint.Quantity)