MNT: simplify and add tests for cbook._check_shape

tacaswell · tacaswell · commit 21082abe0c48 · 2020-05-13T20:23:50.000-04:00
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -2197,24 +2197,20 @@ def _check_shape(_shape, **kwargs):
     for k, v in kwargs.items():
         data_shape = v.shape
 
-        if not (
-            len(data_shape) == len(target_shape)
-            and all(
-                (t == s if t is not None else True)
+        if len(target_shape) != len(data_shape) or any(
+                t not in [s, None]
                 for t, s in zip(target_shape, data_shape)
-            )
         ):
-            def format_dims(target_shape):
-                dim_labels = iter(itertools.chain(
-                    'MNLIJKLH',
-                    (f"D{i}" for i in itertools.count())))
-                text_shape = tuple(n if n is not None else next(dim_labels)
-                                   for n in target_shape)
-                return '(' + ", ".join(str(_) for _ in text_shape) + ')'
+            dim_labels = iter(itertools.chain(
+                'MNLIJKLH',
+                (f"D{i}" for i in itertools.count())))
+            text_shape = ", ".join(str(_) for _ in
+                                   (n if n is not None else next(dim_labels)
+                                    for n in target_shape))
 
             raise ValueError(
                 f"{k!r} must be {len(target_shape)}D "
-                f"with shape {format_dims(target_shape)}. "
+                f"with shape ({text_shape}). "
                 f"Your input has shape {v.shape}."
             )
 
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -1,5 +1,7 @@
 import itertools
 import pickle
+import re
+
 from weakref import ref
 from unittest.mock import patch, Mock
 
@@ -633,3 +635,19 @@ def divisors(n):
         for rstride, cstride in itertools.product(divisors(rows - 1),
                                                   divisors(cols - 1)):
             check(x, rstride=rstride, cstride=cstride)
+
+
+@pytest.mark.parametrize('target,test_shape',
+                         [((None, ), (1, 3)),
+                          ((None, 3), (1,)),
+                          ((None, 3), (1, 2)),
+                          ((1, 5), (1, 9)),
+                          ((None, 2, None), (1, 3, 1))
+                          ])
+def test_check_shape(target, test_shape):
+    error_pattern = (f"^'aardvark' must be {len(target)}D.*" +
+                     re.escape(f'has shape {test_shape}'))
+    data = np.zeros(test_shape)
+    with pytest.raises(ValueError,
+                       match=error_pattern):
+        cbook._check_shape(target, aardvark=data)
