Add pint arrays support

cupy_xarray/accessors.py

@@ -7,8 +7,8 @@
 )
 from xarray.core.pycompat import DuckArrayModule
 
-dsk = DuckArrayModule("dask")
-dask_array_type = dsk.type
+dask_array_type = DuckArrayModule("dask").type
+pint_array_type = DuckArrayModule("pint").type
 
 
 @register_dataarray_accessor("cupy")
@@ -24,9 +24,15 @@ def __init__(self, da):
     @property
     def is_cupy(self):
         """bool: The underlying data is a cupy array."""
-        if isinstance(self.da.data, dask_array_type):
-            return isinstance(self.da.data._meta, cp.ndarray)
-        return isinstance(self.da.data, cp.ndarray)
+        return self._get_datatype(self.da.data)
+
+    @classmethod
+    def _get_datatype(cls, data):
+        if isinstance(data, dask_array_type):
+            return isinstance(data._meta, cp.ndarray)
+        elif isinstance(data, pint_array_type):
+            return cls._get_datatype(data.magnitude)
+        return isinstance(data, cp.ndarray)
 
     def as_cupy(self):
         """
@@ -51,21 +57,20 @@ def as_cupy(self):
         <class 'cupy.core.core.ndarray'>
 
         """
-        if isinstance(self.da.data, dask_array_type):
-            return DataArray(
-                data=self.da.data.map_blocks(cp.asarray),
-                coords=self.da.coords,
-                dims=self.da.dims,
-                name=self.da.name,
-                attrs=self.da.attrs,
+        return self._as_dataarray(self._as_cupy_data(self.da.data))
+
+    @classmethod
+    def _as_cupy_data(cls, data):
+        if isinstance(data, dask_array_type):
+            return data.map_blocks(cp.asarray)
+        if isinstance(data, pint_array_type):
+            from pint import Quantity
+
+            return Quantity(
+                cls._as_cupy_data(data.magnitude),
+                units=data.units,
             )
-        return DataArray(
-            data=cp.asarray(self.da.data),
-            coords=self.da.coords,
-            dims=self.da.dims,
-            name=self.da.name,
-            attrs=self.da.attrs,
-        )
+        return cp.asarray(data)
 
     def as_numpy(self):
         """
@@ -77,29 +82,35 @@ def as_numpy(self):
             DataArray with underlying data cast to numpy.
 
         """
-        if self.is_cupy:
-            if isinstance(self.da.data, dask_array_type):
-                return DataArray(
-                    data=self.da.data.map_blocks(
-                        lambda block: block.get(), dtype=self.da.data._meta.dtype
-                    ),
-                    coords=self.da.coords,
-                    dims=self.da.dims,
-                    name=self.da.name,
-                    attrs=self.da.attrs,
-                )
-            return DataArray(
-                data=self.da.data.get(),
-                coords=self.da.coords,
-                dims=self.da.dims,
-                name=self.da.name,
-                attrs=self.da.attrs,
+        return self._as_dataarray(self._as_numpy_data(self.da.data))
+
+    @classmethod
+    def _as_numpy_data(cls, data):
+        if isinstance(data, dask_array_type):
+            return data.map_blocks(lambda block: block.get(), dtype=data._meta.dtype)
+        if isinstance(data, pint_array_type):
+            from pint import Quantity
+
+            return Quantity(
+                cls._as_numpy_data(data.magnitude),
+                units=data.units,
             )
-        return self.da.as_numpy()
+        if isinstance(data, cp.ndarray):
+            return data.get()
+        return data
 
     def get(self):
         return self.da.data.get()
 
+    def _as_dataarray(self, data):
+        return DataArray(
+            data=data,
+            coords=self.da.coords,
+            dims=self.da.dims,
+            name=self.da.name,
+            attrs=self.da.attrs,
+        )
+
 
 @register_dataset_accessor("cupy")
 class CupyDatasetAccessor:

cupy_xarray/tests/test_accessors.py

@@ -1,10 +1,18 @@
 import numpy as np
 import pytest
 import xarray as xr
-from xarray.core.pycompat import dask_array_type
+from xarray.core.pycompat import DuckArrayModule
+from xarray.tests import requires_pint
+
+dask_array_type = DuckArrayModule("dask").type
+pint_array_type = DuckArrayModule("pint").type
+
+from pint import UnitRegistry
 
 import cupy_xarray  # noqa: F401
 
+ureg = UnitRegistry()
+
 
 @pytest.fixture
 def tutorial_ds_air():
@@ -26,6 +34,28 @@ def tutorial_da_air_dask(tutorial_ds_air_dask):
     return tutorial_ds_air_dask.air
 
 
+@pytest.fixture
+def tutorial_ds_air_pint():
+    return xr.tutorial.load_dataset("air_temperature") * ureg.Quantity("degree_Kelvin")
+
+
+@pytest.fixture
+def tutorial_da_air_pint(tutorial_ds_air_pint):
+    return tutorial_ds_air_pint.air
+
+
+@pytest.fixture
+def tutorial_ds_air_pint_dask():
+    return xr.tutorial.open_dataset(
+        "air_temperature", chunks={"lat": 25, "lon": 25, "time": -1}
+    ) * ureg.Quantity("degree_Kelvin")
+
+
+@pytest.fixture
+def tutorial_da_air_pint_dask(tutorial_ds_air_pint_dask):
+    return tutorial_ds_air_pint_dask.air
+
+
 def test_data_set_accessor(tutorial_ds_air):
     ds = tutorial_ds_air
     assert hasattr(ds, "cupy")
@@ -64,3 +94,35 @@ def test_data_array_accessor_dask(tutorial_da_air_dask):
 
     da = da.cupy.as_numpy()
     assert not da.cupy.is_cupy
+
+
+@requires_pint
+def test_data_array_accessor_pint(tutorial_da_air_pint):
+    da = tutorial_da_air_pint
+    assert hasattr(da, "cupy")
+    assert not da.cupy.is_cupy
+
+    da = da.as_cupy()
+    assert da.cupy.is_cupy
+    assert isinstance(da.data, pint_array_type)
+
+    da = da.cupy.as_numpy()
+    assert not da.cupy.is_cupy
+    assert isinstance(da.data, pint_array_type)
+
+
+@requires_pint
+def test_data_array_accessor_pint_dask(tutorial_da_air_pint_dask):
+    da = tutorial_da_air_pint_dask
+    assert hasattr(da, "cupy")
+    assert not da.cupy.is_cupy
+
+    da = da.as_cupy()
+    assert da.cupy.is_cupy
+    assert isinstance(da.data, pint_array_type)
+    assert isinstance(da.data.magnitude, dask_array_type)
+
+    da = da.cupy.as_numpy()
+    assert not da.cupy.is_cupy
+    assert isinstance(da.data, pint_array_type)
+    assert isinstance(da.data.magnitude, dask_array_type)