ENH: Add argument size validation to quiver. (#14424)

efiring · ImportanceOfBeingErnest · commit e9677c1c9845 · 2019-08-30T14:42:02.000+02:00
Closes #14060.
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -72,12 +72,16 @@
 U, V : 1D or 2D array-like
     The x and y direction components of the arrow vectors.
 
+    They must have the same number of elements, matching the number of arrow
+    locations. They may be masked.
+
 C : 1D or 2D array-like, optional
     Numeric data that defines the arrow colors by colormapping via *norm* and
     *cmap*.
 
     This does not support explicit colors. If you want to set colors directly,
-    use *color* instead.
+    use *color* instead.  The size of *C* must match the number of arrow
+    locations.
 
 units : {'width', 'height', 'dots', 'inches', 'x', 'y' 'xy'}, default: 'width'
     The arrow dimensions (except for *length*) are measured in multiples of
@@ -428,10 +432,13 @@ def _parse_args(*args, caller_name='function'):
         Y = Y.ravel()
         if len(X) == nc and len(Y) == nr:
             X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
+        elif len(X) != len(Y):
+            raise ValueError('X and Y must be the same size, but '
+                             f'X.size is {X.size} and Y.size is {Y.size}.')
     else:
         indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
         X, Y = [np.ravel(a) for a in indexgrid]
-
+    # Size validation for U, V, C is left to the set_UVC method.
     return X, Y, U, V, C
 
 
@@ -589,9 +596,16 @@ def set_UVC(self, U, V, C=None):
         # to an array that might change before draw().
         U = ma.masked_invalid(U, copy=True).ravel()
         V = ma.masked_invalid(V, copy=True).ravel()
-        mask = ma.mask_or(U.mask, V.mask, copy=False, shrink=True)
         if C is not None:
             C = ma.masked_invalid(C, copy=True).ravel()
+        for name, var in zip(('U', 'V', 'C'), (U, V, C)):
+            if var is not None and var.size != self.N:
+                raise ValueError(f'Argument {name} has a size {var.size}'
+                                 f' which does not match {self.N},'
+                                 ' the number of arrow positions')
+
+        mask = ma.mask_or(U.mask, V.mask, copy=False, shrink=True)
+        if C is not None:
             mask = ma.mask_or(mask, C.mask, copy=False, shrink=True)
             if mask is ma.nomask:
                 C = C.filled()
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -51,6 +51,27 @@ def test_quiver_number_of_args():
         plt.quiver(X, X, X, X, X, X)
 
 
+def test_quiver_arg_sizes():
+    X2 = [1, 2]
+    X3 = [1, 2, 3]
+    with pytest.raises(ValueError,
+            match=('X and Y must be the same size, but '
+                   'X.size is 2 and Y.size is 3.')):
+        plt.quiver(X2, X3, X2, X2)
+    with pytest.raises(ValueError,
+            match=('Argument U has a size 3 which does not match 2,'
+                   ' the number of arrow positions')):
+        plt.quiver(X2, X2, X3, X2)
+    with pytest.raises(ValueError,
+            match=('Argument V has a size 3 which does not match 2,'
+                   ' the number of arrow positions')):
+        plt.quiver(X2, X2, X2, X3)
+    with pytest.raises(ValueError,
+            match=('Argument C has a size 3 which does not match 2,'
+                   ' the number of arrow positions')):
+        plt.quiver(X2, X2, X2, X2, X3)
+
+
 def test_no_warnings():
     fig, ax = plt.subplots()
 
