BUG: Gracefully handle non-finite z in tricontour (matplotlib#10167)

ianthomas23 · ianthomas23 · commit fdbde0ce21e6 · 2019-04-25T14:39:37.000+01:00
diff --git a/lib/matplotlib/tests/test_triangulation.py b/lib/matplotlib/tests/test_triangulation.py
@@ -1115,3 +1115,27 @@ def test_qhull_large_offset():
     triang = mtri.Triangulation(x, y)
     triang_offset = mtri.Triangulation(x + offset, y + offset)
     assert len(triang.triangles) == len(triang_offset.triangles)
+
+
+def test_tricontour_non_finite_z():
+    # github issue 10167.
+    x = [0, 1, 0, 1]
+    y = [0, 0, 1, 1]
+    triang = mtri.Triangulation(x, y)
+    plt.figure()
+
+    with pytest.raises(ValueError) as excinfo:
+        plt.tricontourf(triang, [0, 1, 2, np.inf])
+    excinfo.match(r'z array cannot contain non-finite values')
+
+    with pytest.raises(ValueError) as excinfo:
+        plt.tricontourf(triang, [0, 1, 2, -np.inf])
+    excinfo.match(r'z array cannot contain non-finite values')
+
+    with pytest.raises(ValueError) as excinfo:
+        plt.tricontourf(triang, [0, 1, 2, np.nan])
+    excinfo.match(r'z array cannot contain non-finite values')
+
+    with pytest.raises(ValueError) as excinfo:
+        plt.tricontourf(triang, np.ma.array([0, 1, 2, 3], mask=[1, 0, 0, 0]))
+    excinfo.match(r'z cannot be a masked array')
diff --git a/lib/matplotlib/tri/tricontour.py b/lib/matplotlib/tri/tricontour.py
@@ -79,12 +79,22 @@ def _contour_args(self, args, kwargs):
             fn = 'contour'
         tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args,
                                                                    **kwargs)
-        z = np.asarray(args[0])
+        z = np.ma.asarray(args[0])
         if z.shape != tri.x.shape:
             raise ValueError('z array must have same length as triangulation x'
                              ' and y arrays')
-        self.zmax = z.max()
-        self.zmin = z.min()
+
+        # z values must be finite, only need to check points that are included
+        # in the triangulation.
+        z_check = z[np.unique(tri.get_masked_triangles())]
+        if np.ma.is_masked(z_check):
+            raise ValueError('z cannot be a masked array')
+        if not np.isfinite(z_check).all():
+            raise ValueError('z array cannot contain non-finite values')
+
+        z = np.ma.masked_invalid(z, copy=False)
+        self.zmax = float(z.max())
+        self.zmin = float(z.min())
         if self.logscale and self.zmin <= 0:
             raise ValueError('Cannot %s log of negative values.' % fn)
         self._contour_level_args(z, args[1:])
