Check number of positional arguments passed to quiver()

timhoffm · timhoffm · commit 370cc5b5d210 · 2019-04-29T16:27:50.000+02:00
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -2148,16 +2148,6 @@ def _check_and_log_subprocess(command, logger, **kwargs):
     return report
 
 
-def _check_not_matrix(**kwargs):
-    """
-    If any value in *kwargs* is a `np.matrix`, raise a TypeError with the key
-    name in its message.
-    """
-    for k, v in kwargs.items():
-        if isinstance(v, np.matrix):
-            raise TypeError(f"Argument {k!r} cannot be a np.matrix")
-
-
 def _check_in_list(values, **kwargs):
     """
     For each *key, value* pair in *kwargs*, check that *value* is in *values*;
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -360,31 +360,53 @@ def quiverkey_doc(self):
         return self.__init__.__doc__
 
 
-# This is a helper function that parses out the various combination of
-# arguments for doing colored vector plots.  Pulling it out here
-# allows both Quiver and Barbs to use it
-def _parse_args(*args):
-    X = Y = U = V = C = None
-    args = list(args)
-
-    # The use of atleast_1d allows for handling scalar arguments while also
-    # keeping masked arrays
-    if len(args) == 3 or len(args) == 5:
-        C = np.atleast_1d(args.pop(-1))
-    V = np.atleast_1d(args.pop(-1))
-    U = np.atleast_1d(args.pop(-1))
-    cbook._check_not_matrix(U=U, V=V, C=C)
-    if U.ndim == 1:
-        nr, nc = 1, U.shape[0]
+def _parse_args(*args, caller_name='function'):
+    """
+    Helper function to parse positional parameters for colored vector plots.
+
+    This is currently used for Quiver and Barbs.
+
+    Parameters
+    ----------
+    *args : list
+        list of 2-5 arguments. Depending on their number they are parsed to::
+
+            U, V
+            U, V, C
+            X, Y, U, V
+            X, Y, U, V, C
+
+    caller_name : str
+        Name of the calling method (used in error messages).
+    """
+    X = Y = C = None
+
+    len_args = len(args)
+    if len_args == 2:
+        # The use of atleast_1d allows for handling scalar arguments while also
+        # keeping masked arrays
+        U, V = np.atleast_1d(*args)
+    elif len_args == 3:
+        U, V, C = np.atleast_1d(*args)
+    elif len_args == 4:
+        X, Y, U, V = np.atleast_1d(*args)
+    elif len_args == 5:
+        X, Y, U, V, C = np.atleast_1d(*args)
     else:
-        nr, nc = U.shape
-    if len(args) == 2:  # remaining after removing U,V,C
-        X, Y = [np.array(a).ravel() for a in args]
+        raise TypeError(f'{caller_name} takes 2-5 positional arguments but '
+                        f'{len_args} were given')
+
+    nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape
+
+    if X is not None:
+        X = X.ravel()
+        Y = Y.ravel()
         if len(X) == nc and len(Y) == nr:
             X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
     else:
         indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
         X, Y = [np.ravel(a) for a in indexgrid]
+
     return X, Y, U, V, C
 
 
@@ -426,7 +448,7 @@ def __init__(self, ax, *args,
         %s
         """
         self.ax = ax
-        X, Y, U, V, C = _parse_args(*args)
+        X, Y, U, V, C = _parse_args(*args, caller_name='quiver()')
         self.X = X
         self.Y = Y
         self.XY = np.column_stack((X, Y))
@@ -941,7 +963,7 @@ def __init__(self, ax, *args,
             kw['linewidth'] = 1
 
         # Parse out the data arrays from the various configurations supported
-        x, y, u, v, c = _parse_args(*args)
+        x, y, u, v, c = _parse_args(*args, caller_name='barbs()')
         self.x = x
         self.y = y
         xy = np.column_stack((x, y))
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -41,6 +41,16 @@ def test_quiver_key_memory_leak():
     assert sys.getrefcount(qk) == 2
 
 
+def test_quiver_number_of_args():
+    X = [1, 2]
+    with pytest.raises(TypeError,
+            match='takes 2-5 positional arguments but 1 were given'):
+        plt.quiver(X)
+    with pytest.raises(TypeError,
+            match='takes 2-5 positional arguments but 6 were given'):
+        plt.quiver(X, X, X, X, X, X)
+
+
 def test_no_warnings():
     fig, ax = plt.subplots()
 
