allow np-array levels and colors in 2D plots (#3295)

* test if levels is None

* allow np levels and color list

* whats-new

* Update doc/whats-new.rst

Co-Authored-By: Deepak Cherian <dcherian@users.noreply.github.com>

doc/whats-new.rst

@@ -165,6 +165,8 @@ Bug fixes
   By `Hasan Ahmad <https://github.com/HasanAhmadQ7>`_.
 - Prevent :py:meth:`~xarray.DataArray.argmax` and :py:meth:`~xarray.DataArray.argmin` from calling
   dask compute (:issue:`3237`). By `Ulrich Herter <https://github.com/ulijh>`_.
+- Plots in 2 dimensions (pcolormesh, contour) now allow to specify levels as numpy
+  array (:issue:`3284`). By `Mathias Hauser <https://github.com/mathause>`_.
 
 .. _whats-new.0.12.3:
 

xarray/plot/utils.py

@@ -731,7 +731,7 @@ def _process_cmap_cbar_kwargs(
 
     # colors is only valid when levels is supplied or the plot is of type
     # contour or contourf
-    if colors and (("contour" not in func.__name__) and (not levels)):
+    if colors and (("contour" not in func.__name__) and (levels is None)):
         raise ValueError("Can only specify colors with contour or levels")
 
     # we should not be getting a list of colors in cmap anymore

xarray/tests/test_plot.py

@@ -1283,26 +1283,38 @@ class TestContour(Common2dMixin, PlotTestCase):
 
     plotfunc = staticmethod(xplt.contour)
 
+    # matplotlib cmap.colors gives an rgbA ndarray
+    # when seaborn is used, instead we get an rgb tuple
+    @staticmethod
+    def _color_as_tuple(c):
+        return tuple(c[:3])
+
     def test_colors(self):
-        # matplotlib cmap.colors gives an rgbA ndarray
-        # when seaborn is used, instead we get an rgb tuple
-        def _color_as_tuple(c):
-            return tuple(c[:3])
 
         # with single color, we don't want rgb array
         artist = self.plotmethod(colors="k")
         assert artist.cmap.colors[0] == "k"
 
         artist = self.plotmethod(colors=["k", "b"])
-        assert _color_as_tuple(artist.cmap.colors[1]) == (0.0, 0.0, 1.0)
+        assert self._color_as_tuple(artist.cmap.colors[1]) == (0.0, 0.0, 1.0)
 
         artist = self.darray.plot.contour(
             levels=[-0.5, 0.0, 0.5, 1.0], colors=["k", "r", "w", "b"]
         )
-        assert _color_as_tuple(artist.cmap.colors[1]) == (1.0, 0.0, 0.0)
-        assert _color_as_tuple(artist.cmap.colors[2]) == (1.0, 1.0, 1.0)
+        assert self._color_as_tuple(artist.cmap.colors[1]) == (1.0, 0.0, 0.0)
+        assert self._color_as_tuple(artist.cmap.colors[2]) == (1.0, 1.0, 1.0)
+        # the last color is now under "over"
+        assert self._color_as_tuple(artist.cmap._rgba_over) == (0.0, 0.0, 1.0)
+
+    def test_colors_np_levels(self):
+
+        # https://github.com/pydata/xarray/issues/3284
+        levels = np.array([-0.5, 0.0, 0.5, 1.0])
+        artist = self.darray.plot.contour(levels=levels, colors=["k", "r", "w", "b"])
+        assert self._color_as_tuple(artist.cmap.colors[1]) == (1.0, 0.0, 0.0)
+        assert self._color_as_tuple(artist.cmap.colors[2]) == (1.0, 1.0, 1.0)
         # the last color is now under "over"
-        assert _color_as_tuple(artist.cmap._rgba_over) == (0.0, 0.0, 1.0)
+        assert self._color_as_tuple(artist.cmap._rgba_over) == (0.0, 0.0, 1.0)
 
     def test_cmap_and_color_both(self):
         with pytest.raises(ValueError):