Merge pull request #7570 from dstansby/scatter-nan-colors

tacaswell · tacaswell · commit e09a6f74873d · 2016-12-11T00:00:32.000-05:00
MNT: Correctly skip colors for nan points given to scatter

Conflicts:
    lib/matplotlib/tests/test_axes.py
       Conflicts due to tests added on both branches
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3986,8 +3986,11 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
         else:
             colors = None  # use cmap, norm after collection is created
 
-        # c will be unchanged unless it is the same length as x:
-        x, y, s, c = cbook.delete_masked_points(x, y, s, c)
+        # Anything in maskargs will be unchanged unless it is the same length
+        # as x:
+        maskargs = x, y, s, c, colors, edgecolors, linewidths
+        x, y, s, c, colors, edgecolors, linewidths =\
+            cbook.delete_masked_points(*maskargs)
 
         scales = s   # Renamed for readability below.
 
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -4785,6 +4785,22 @@ def test_log_margins():
     assert_allclose([xlim0t + delta, xlim1t - delta], [x0t, x1t])
 
 
+@cleanup
+def test_scatter_color_masking():
+    x = np.array([1, 2, 3])
+    y = np.array([1, np.nan, 3])
+    colors = np.array(['k', 'w', 'k'])
+    linewidths = np.array([1, 2, 3])
+    s = plt.scatter(x, y, color=colors, linewidths=linewidths)
+
+    facecolors = s.get_facecolors()
+    linecolors = s.get_edgecolors()
+    linewidths = s.get_linewidths()
+    assert_array_equal(facecolors[1], np.array([0, 0, 0, 1]))
+    assert_array_equal(linecolors[1], np.array([0, 0, 0, 1]))
+    assert linewidths[1] == 3
+
+
 if __name__ == '__main__':
     import nose
     import sys
