Plots get labels from pint arrays

xarray/core/pycompat.py

@@ -45,3 +45,13 @@ def is_duck_dask_array(x):
 except ImportError:  # pragma: no cover
     cupy_version = LooseVersion("0.0.0")
     cupy_array_type = ()
+
+try:
+    # solely for isinstance checks
+    import pint
+
+    pint_version = LooseVersion(pint.__version__)
+    pint_array_type = (pint.Quantity,)
+except ImportError:  # pragma: no cover
+    pint_version = LooseVersion("0.0.0")
+    pint_array_type = ()

xarray/core/variable.py

@@ -33,6 +33,7 @@
     dask_array_type,
     integer_types,
     is_duck_dask_array,
+    pint_array_type,
 )
 from .utils import (
     NdimSizeLenMixin,
@@ -252,19 +253,28 @@ def _as_array_or_item(data):
     Importantly, this function does not copy data if it is already an ndarray -
     otherwise, it will not be possible to update Variable values in place.
 
-    This function mostly exists because 0-dimensional ndarrays with
-    dtype=datetime64 are broken :(
-    https://github.com/numpy/numpy/issues/4337
-    https://github.com/numpy/numpy/issues/7619
-
-    TODO: remove this (replace with np.asarray) once these issues are fixed
     """
-    data = data.get() if isinstance(data, cupy_array_type) else np.asarray(data)
+
+    # TODO replace this with an entrypoint to allow duck-array libraries to
+    # specify how they want to return their values?
+    if isinstance(data, cupy_array_type):
+        data = data.get()
+    elif isinstance(data, pint_array_type):
+        data = data.magnitude
+    else:
+        data = np.asarray(data)
+
+    # Exists because 0-dimensional ndarrays with
+    # dtype=datetime64 are broken :(
+    # https://github.com/numpy/numpy/issues/4337
+    # https://github.com/numpy/numpy/issues/7619
+    # TODO: remove this (replace with np.asarray) once these issues are fixed
     if data.ndim == 0:
         if data.dtype.kind == "M":
             data = np.datetime64(data, "ns")
         elif data.dtype.kind == "m":
             data = np.timedelta64(data, "ns")
+
     return data
 
 

xarray/plot/utils.py

@@ -9,6 +9,7 @@
 import pandas as pd
 
 from ..core.options import OPTIONS
+from ..core.pycompat import pint_array_type
 from ..core.utils import is_scalar
 
 try:
@@ -474,12 +475,19 @@ def label_from_attrs(da, extra=""):
     else:
         name = ""
 
-    if da.attrs.get("units"):
-        units = " [{}]".format(da.attrs["units"])
-    elif da.attrs.get("unit"):
-        units = " [{}]".format(da.attrs["unit"])
+    def _get_units_from_attrs(da):
+        if da.attrs.get("units"):
+            units = " [{}]".format(da.attrs["units"])
+        elif da.attrs.get("unit"):
+            units = " [{}]".format(da.attrs["unit"])
+        else:
+            units = ""
+        return units
+
+    if isinstance(da.data, pint_array_type):
+        units = " [{}]".format(str(da.data.units))
     else:
-        units = ""
+        units = _get_units_from_attrs(da)
 
     return "\n".join(textwrap.wrap(name + extra + units, 30))
 

xarray/tests/test_units.py

@@ -5,10 +5,22 @@
 import pandas as pd
 import pytest
 
+try:
+    import matplotlib.pyplot as plt
+except ImportError:
+    pass
+
 import xarray as xr
 from xarray.core import dtypes, duck_array_ops
 
-from . import assert_allclose, assert_duckarray_allclose, assert_equal, assert_identical
+from . import (
+    assert_allclose,
+    assert_duckarray_allclose,
+    assert_equal,
+    assert_identical,
+    requires_matplotlib,
+)
+from .test_plot import PlotTestCase
 from .test_variable import _PAD_XR_NP_ARGS
 
 pint = pytest.importorskip("pint")
@@ -5564,3 +5576,29 @@ def test_merge(self, variant, unit, error, dtype):
 
         assert_units_equal(expected, actual)
         assert_equal(expected, actual)
+
+
+@requires_matplotlib
+class TestPlots(PlotTestCase):
+    def test_units_in_line_plot_labels(self):
+        arr = np.linspace(1, 10, 3) * unit_registry.Pa
+        # TODO make coord a Quantity once unit-aware indexes supported
+        x_coord = xr.DataArray(
+            np.linspace(1, 3, 3), dims="x", attrs={"units": "meters"}
+        )
+        da = xr.DataArray(data=arr, dims="x", coords={"x": x_coord}, name="pressure")
+
+        da.plot.line()
+
+        ax = plt.gca()
+        assert ax.get_ylabel() == "pressure [pascal]"
+        assert ax.get_xlabel() == "x [meters]"
+
+    def test_units_in_2d_plot_labels(self):
+        arr = np.ones((2, 3)) * unit_registry.Pa
+        da = xr.DataArray(data=arr, dims=["x", "y"], name="pressure")
+
+        fig, (ax, cax) = plt.subplots(1, 2)
+        ax = da.plot.contourf(ax=ax, cbar_ax=cax, add_colorbar=True)
+
+        assert cax.get_ylabel() == "pressure [pascal]"