Quadmesh.set_array validates dimensions

jeffreypaul15 · jeffreypaul15 · commit 4554737296c3 · 2021-05-12T22:39:24.000+05:30
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -367,6 +367,20 @@ def set_array(self, A):
         ----------
         A : ndarray or None
         """
+        if isinstance(self, mpl.collections.QuadMesh):
+            A_ = A.ravel()
+            if self._shading in ['flat', 'nearest', 'auto']:
+                if not len(A_) == self._meshWidth * self._meshHeight:
+                    raise TypeError(
+                        'Dimensions of A %s are incompatible with '
+                        'X (%d) and/or Y (%d)' %
+                        (A.shape, self._meshWidth, self._meshHeight))
+            elif self._shading == 'gouraud':
+                if not len(A_) == (self._meshHeight + 1) * (self._meshWidth + 1):
+                    raise TypeError(
+                        'Dimensions of A %s are incompatible with '
+                        'X (%d) and/or Y (%d)' %
+                        (A.shape, self._meshWidth, self._meshHeight))
         self._A = A
 
     def get_array(self):
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -698,6 +698,35 @@ def test_singleton_autolim():
     np.testing.assert_allclose(ax.get_xlim(), [-0.06, 0.06])
 
 
+def test_quadmesh_set_array_validation():
+    x = np.arange(4)
+    y = np.arange(4)
+    z = np.arange(9).reshape((3, 3))
+    fig, ax = plt.subplots()
+    coll = ax.pcolormesh(x, y, np.ones(z.shape))
+    z = np.arange(16).reshape((4, 4))
+
+    with pytest.raises(ValueError, match="Dimensions of A (4, 4) are incompatible with X (3) and/or Y (3)"):
+        coll.set_array(z)
+    with pytest.raises(ValueError, match="Dimensions of A (16,) are incompatible with X (3) and/or Y (3)"):
+        coll.set_array(z.ravel())
+    
+    fig, ax = plt.subplots()
+    coll = ax.pcolormesh(x, y, np.ones(z.shape), shading='nearest')
+    with pytest.raises(ValueError, match="Dimensions of A (3, 3) are incompatible with X (4) and/or Y (4)"):
+        z = np.arange(9).reshape((3, 3))
+        coll.set_array(z)
+
+    x = np.arange(3)
+    y = np.arange(3)
+    z = np.arange(9).reshape((3, 3))
+    fig, ax = plt.subplots()
+    coll = ax.pcolormesh(x, y, np.ones(z.shape), shading='gouraud')
+    z = np.arange(16).reshape((4, 4))
+    with pytest.raises(ValueError, match="Dimensions of A (4, 4) are incompatible with X (2) and/or Y (2)"):
+        coll.set_array(z)
+
+
 def test_quadmesh_set_array():
     x = np.arange(4)
     y = np.arange(4)
