Refactor c_array checking; Lower N in color len mismatch test

bcongdon · bcongdon · commit d1bf75fa2047 · 2016-10-31T06:42:16.000-05:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3992,6 +3992,7 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
 
         # np.ma.ravel yields an ndarray, not a masked array,
         # unless its argument is a masked array.
+        xy_shape = (np.shape(x), np.shape(y))
         x = np.ma.ravel(x)
         y = np.ma.ravel(y)
         if x.size != y.size:
@@ -4014,18 +4015,22 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
         else:
             try:
                 c_array = np.asanyarray(c, dtype=float)
+                if c_array.shape in xy_shape:
+                    c = np.ma.ravel(c_array)
+                else:
+                    # Wrong size; it must not be intended for mapping.
+                    c_array = None
             except ValueError:
                 # Failed to make a floating-point array; c must be color specs.
                 c_array = None
-            else:
-                if c_array.size == x.size:
-                    c = np.ma.ravel(c_array)
-                elif c_array.size not in (3, 4):
-                    # Wrong size. Not a rgb/rgba and not same size as x
-                    raise ValueError("x and c must be the same size")
 
         if c_array is None:
             colors = c     # must be acceptable as PathCollection facecolors
+            try:
+                mcolors.to_rgba_array(colors)
+            except ValueError:
+                # c not acceptable as PathCollection facecolor
+                raise ValueError("c not acceptable as color sequence")
         else:
             colors = None  # use cmap, norm after collection is created
 
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -4717,9 +4717,12 @@ def test_fillbetween_cycle():
 
 @cleanup
 def test_color_length_mismatch():
-    N = 500
-    x, y = np.random.rand(N), np.random.rand(N)
-    colors = np.random.rand(N+1)
+    N = 5
+    x, y = np.arange(N), np.arange(N)
+    colors = np.arange(N+1)
     fig, ax = plt.subplots()
     with pytest.raises(ValueError):
         ax.scatter(x, y, c=colors)
+    c_rgb = (0.5, 0.5, 0.5)
+    ax.scatter(x, y, c=c_rgb)
+    ax.scatter(x, y, c=[c_rgb] * N)
