From 370cc5b5d210af3e7034d77e0d54a44e941134af Mon Sep 17 00:00:00 2001 From: Tim Hoffmann <2836374+timhoffm@users.noreply.github.com> Date: Mon, 29 Apr 2019 16:27:50 +0200 Subject: [PATCH] Check number of positional arguments passed to quiver() --- lib/matplotlib/cbook/__init__.py | 10 ----- lib/matplotlib/quiver.py | 64 +++++++++++++++++++---------- lib/matplotlib/tests/test_quiver.py | 10 +++++ 3 files changed, 53 insertions(+), 31 deletions(-) diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py index bbeb82b18f1a..c762e7318ea6 100644 --- 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 index 6bb3226c9bf3..e3b3257e7286 100644 --- 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 index e93eaf87f5e6..de17c1f5bcb9 100644 --- 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()