Resolve inconsistent NaN handling behaviour

henryhu123 · DennisTismenko · commit db133b0696b6 · 2020-03-13T01:42:13.000-04:00
Combine array masks rather than deleting masked points to maintain consistency across the project. Add appropriate test cases for validating color correctness for hlines and vlines. Fixes issue #13799.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -1122,15 +1122,17 @@ def hlines(self, y, xmin, xmax, colors='k', linestyles='solid',
         if not np.iterable(xmax):
             xmax = [xmax]
 
-        y, xmin, xmax = cbook.delete_masked_points(y, xmin, xmax)
-
+        # Create and combine masked_arrays from input
+        y, xmin, xmax = cbook._combine_masks(y, xmin, xmax)
         y = np.ravel(y)
-        xmin = np.resize(xmin, y.shape)
-        xmax = np.resize(xmax, y.shape)
-
-        verts = [((thisxmin, thisy), (thisxmax, thisy))
-                 for thisxmin, thisxmax, thisy in zip(xmin, xmax, y)]
-        lines = mcoll.LineCollection(verts, colors=colors,
+        xmin = np.ma.resize(xmin, y.shape)
+        xmax = np.ma.resize(xmax, y.shape)
+
+        masked_verts = [np.ma.array([y_xmin, y_xmax])
+                        for y_xmin, y_xmax in
+                        zip(np.ma.array([xmin, y]).T,
+                            np.ma.array([xmax, y]).T)]
+        lines = mcoll.LineCollection(masked_verts, colors=colors,
                                      linestyles=linestyles, label=label)
         self.add_collection(lines, autolim=False)
         lines.update(kwargs)
@@ -1200,15 +1202,17 @@ def vlines(self, x, ymin, ymax, colors='k', linestyles='solid',
         if not np.iterable(ymax):
             ymax = [ymax]
 
-        x, ymin, ymax = cbook.delete_masked_points(x, ymin, ymax)
-
+        # Create and combine masked_arrays from input
+        x, ymin, ymax = cbook._combine_masks(x, ymin, ymax)
         x = np.ravel(x)
-        ymin = np.resize(ymin, x.shape)
-        ymax = np.resize(ymax, x.shape)
-
-        verts = [((thisx, thisymin), (thisx, thisymax))
-                 for thisx, thisymin, thisymax in zip(x, ymin, ymax)]
-        lines = mcoll.LineCollection(verts, colors=colors,
+        ymin = np.ma.resize(ymin, x.shape)
+        ymax = np.ma.resize(ymax, x.shape)
+
+        masked_verts = [np.ma.array([xymin, xymax])
+                        for xymin, xymax in
+                        zip(np.ma.array([x, ymin]).T,
+                            np.ma.array([x, ymax]).T)]
+        lines = mcoll.LineCollection(masked_verts, colors=colors,
                                      linestyles=linestyles, label=label)
         self.add_collection(lines, autolim=False)
         lines.update(kwargs)
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3992,6 +3992,74 @@ def test_hlines():
     ax5.set_ylim(0, 15)
 
 
+def generate_hlines_with_colors_inputs():
+    colors = cycler('colors', [['red', 'green', 'blue', 'purple', 'orange']])
+    base_y = cycler(
+        'y', [[1, 2, 3, np.nan, 5], np.ma.masked_equal([1, 2, 3, 4, 5], 4)])
+    base_max = cycler('xmax', [np.ones(5)])
+    base_min = cycler('xmin', [0])
+    base_line_width = cycler('linewidth', [5])
+    original = base_y * base_min * base_max * colors * base_line_width
+
+    expect = {
+        'y': [1, 2, 3, 5],
+        'xmin': 0, 'xmax': [1, 2, 3, 5],
+        'colors': ['red', 'green', 'blue', 'orange']}
+
+    result = []
+
+    for i in original:
+        result.append({'original': i, 'expect': expect})
+
+    return result
+
+
+@pytest.mark.parametrize('kwargs', generate_hlines_with_colors_inputs())
+@check_figures_equal(extensions=["png"])
+def test_hlines_with_colors(fig_test, fig_ref, kwargs):
+    fig_test, ax0 = plt.subplots()
+    test = kwargs.pop('original')
+    ax0.hlines(**test)
+
+    fig_ref, ax1 = plt.subplots()
+    ref = kwargs.pop('expect')
+    ax1.hlines(**ref)
+
+
+def generate_vlines_with_colors_inputs():
+    colors = cycler('colors', [['red', 'green', 'blue', 'purple', 'orange']])
+    base_x = cycler(
+        'x', [[1, 2, 3, np.nan, 5], np.ma.masked_equal([1, 2, 3, 4, 5], 4)])
+    base_max = cycler('ymax', [np.ones(5)])
+    base_min = cycler('ymin', [0])
+    base_line_width = cycler('linewidth', [5])
+    original = base_x * base_min * base_max * colors * base_line_width
+
+    expect = {
+        'x': [1, 2, 3, 5],
+        'ymin': 0,
+        'ymax': [1, 2, 3, 5],
+        'colors': ['red', 'green', 'blue', 'orange']}
+
+    result = []
+    for i in original:
+        result.append({'original': i, 'expect': expect})
+
+    return result
+
+
+@pytest.mark.parametrize('kwargs', generate_vlines_with_colors_inputs())
+@check_figures_equal(extensions=["png"])
+def test_vlines_with_colors(fig_test, fig_ref, kwargs):
+    fig_test, ax0 = plt.subplots()
+    test = kwargs.pop('original')
+    ax0.vlines(**test)
+
+    fig_ref, ax1 = plt.subplots()
+    ref = kwargs.pop('expect')
+    ax1.vlines(**ref)
+
+
 @image_comparison(['step_linestyle', 'step_linestyle'], remove_text=True)
 def test_step_linestyle():
     x = y = np.arange(10)
