Merge pull request #7363 from bcongdon/scatter-error-msg-fix

efiring · web-flow · commit 5901b3844257 · 2016-11-09T09:04:09.000-10:00
Add appropriate error on color size mismatch in `scatter`
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3952,6 +3952,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:
@@ -3974,7 +3975,7 @@ 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.size == x.size:
+                if c_array.shape in xy_shape:
                     c = np.ma.ravel(c_array)
                 else:
                     # Wrong size; it must not be intended for mapping.
@@ -3984,7 +3985,14 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
                 c_array = None
 
         if c_array is None:
-            colors = c     # must be acceptable as PathCollection facecolors
+            try:
+                # must be acceptable as PathCollection facecolors
+                colors = mcolors.to_rgba_array(c)
+            except ValueError:
+                # c not acceptable as PathCollection facecolor
+                msg = ("c of shape {0} not acceptable as a color sequence "
+                       "for x with size {1}, y with size {2}")
+                raise ValueError(msg.format(c.shape, x.size, y.size))
         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
@@ -4800,15 +4800,13 @@ def test_fillbetween_cycle():
 
 
 @cleanup
-def test_log_margins():
-    plt.rcParams['axes.autolimit_mode'] = 'data'
+def test_color_length_mismatch():
+    N = 5
+    x, y = np.arange(N), np.arange(N)
+    colors = np.arange(N+1)
     fig, ax = plt.subplots()
-    margin = 0.05
-    ax.set_xmargin(margin)
-    ax.semilogx([1, 10], [1, 10])
-    xlim0, xlim1 = ax.get_xlim()
-    transform = ax.xaxis.get_transform()
-    xlim0t, xlim1t = transform.transform([xlim0, xlim1])
-    x0t, x1t = transform.transform([1, 10])
-    delta = (x1t - x0t) * margin
-    assert_allclose([xlim0t + delta, xlim1t - delta], [x0t, x1t])
+    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)
