ENH: Catch masked array and invalid x, y to pcolor/mesh

jklymak · jklymak · commit 4733921e4fe1 · 2017-11-10T14:27:40.000-08:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5182,7 +5182,22 @@ def _pcolorargs(funcname, *args, **kw):
             return X, Y, C
 
         if len(args) == 3:
+            # Check x and y for bad data...
             X, Y, C = [np.asanyarray(a) for a in args]
+            for a in X, Y:
+                a = cbook.safe_masked_invalid(a)
+                if np.ma.is_masked(a):
+                    raise TypeError(
+                        'x and y arguments to %s cannot have '
+                        'non-finite values and cannot be type '
+                        'numpy.ma.core.MaskedArray with masked values'
+                        % funcname)
+            # if it didn't fail the above it is masked but has no
+            # masked values, so strip the mask:
+            if type(X) is np.ma.core.MaskedArray:
+                X = X.data
+            if type(Y) is np.ma.core.MaskedArray:
+                Y = Y.data
             numRows, numCols = C.shape
         else:
             raise TypeError(
@@ -5577,7 +5592,6 @@ def pcolormesh(self, *args, **kwargs):
         # convert to one dimensional arrays
         C = C.ravel()
         coords = np.column_stack((X, Y)).astype(float, copy=False)
-
         collection = mcoll.QuadMesh(Nx - 1, Ny - 1, coords,
                                     antialiased=antialiased, shading=shading,
                                     **kwargs)
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1190,6 +1190,13 @@ def test_pcolorargs():
         ax.pcolormesh(x, y, Z[:-1, :-1], shading="gouraud")
     with pytest.raises(TypeError):
         ax.pcolormesh(X, Y, Z[:-1, :-1], shading="gouraud")
+    with pytest.raises(TypeError):
+        xx = x
+        xx[0] = np.NaN
+        ax.pcolormesh(xx, y, Z[:-1, :-1])
+    with pytest.raises(TypeError):
+        xx = np.ma.masked(x, mask=(x < 0))
+        ax.pcolormesh(xx, y, Z[:-1, :-1])
 
 
 @image_comparison(baseline_images=['canonical'])
