Changing to check_figures_equal

henryhu123 · henryhu123 · commit b87762fd1a2f · 2020-03-10T23:56:23.000-04:00
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3991,89 +3991,76 @@ def test_hlines():
     ax5.set_ylim(0, 15)
 
 
-@image_comparison(
-    ['vlines_colors_basic', 'vlines_colors_with_nan', 'vlines_colors_masked'],
-    extensions=['png'])
-def test_vlines_with_colors():
-    # normal
-    colors = ['red', 'green', 'blue', 'purple', 'orange']
-    x1 = [2, 3, 4, 5, 7]
-    y1 = [2, -6, 3, 8, 2]
-    fig1, ax1 = plt.subplots()
-    ax1.vlines(x1, 0, y1, colors=colors, linewidth=5)
-
-    # GH #7406
-    x2 = [2, 3, 4, 5, 6, 7]
-    y2 = [2, -6, 3, 8, np.nan, 2]
-    fig2, (ax2, ax3, ax4) = plt.subplots(nrows=3, figsize=(4, 8))
-    ax2.vlines(x2, 0, y2, colors=colors, linewidth=5)
-
-    x3 = [2, 3, 4, 5, 6, 7]
-    y3 = [np.nan, 2, -6, 3, 8, 2]
-    ax3.vlines(x3, 0, y3, colors=colors, linewidth=3, linestyle='--')
-
-    x4 = [2, 3, 4, 5, 6, 7]
-    y4 = [np.nan, 2, -6, 3, 8, np.nan]
-    ax4.vlines(x4, 0, y4, colors=colors, linewidth=2)
-
-    # tweak the x-axis so we can see the lines better
-    for ax in [ax1, ax2, ax3, ax4]:
-        ax.set_xlim(0, 10)
-
-    # check that the y-lims are all automatically the same
-    assert ax1.get_ylim() == ax2.get_ylim()
-    assert ax1.get_ylim() == ax3.get_ylim()
-    assert ax1.get_ylim() == ax4.get_ylim()
+@check_figures_equal()
+def test_vlines_with_nan_colors(fig_test, fig_ref):
+    colors1 = ['red', 'green', 'blue', 'purple', 'orange', 'black']
+    x1 = [2, 3, 4, 5, 6, 7]
+    y1 = [2, -6, 3, 8, np.nan, 2]
 
-    fig3, ax5 = plt.subplots()
-    x5 = np.ma.masked_equal([2, 4, 6, 8, 10, 12], 8)
-    ymin5 = np.ma.masked_equal([0, 1, -1, 0, 2, 1], 2)
-    ymax5 = np.ma.masked_equal([13, 14, 15, 16, 17, 18], 18)
-    ax5.vlines(x5, ymin5, ymax5, colors=colors, linewidth=2)
-    ax5.set_xlim(0, 15)
+    fig_test, ax1 = plt.subplots()
+    ax1.vlines(x1, 0, y1, colors=colors1, linewidth=5)
 
+    colors2 = ['red', 'green', 'blue', 'purple', 'black']
+    x2 = [2, 3, 4, 5, 7]
+    y2 = [2, -6, 3, 8, 2]
+    # Reference image
+    fig_ref, ax2 = plt.subplots()
+    ax2.vlines(x2, 0, y2, colors=colors2, linewidth=5)
 
-@image_comparison(
-    ['hlines_colors_basic', 'hlines_colors_with_nan', 'hlines_colors_masked'],
-                  extensions=['png'])
-def test_hlines_with_colors():
 
-    colors = ['red', 'green', 'blue', 'purple', 'orange']
-    # normal
-    y1 = [2, 3, 4, 5, 7]
-    x1 = [2, -6, 3, 8, 2]
-    fig1, ax1 = plt.subplots()
-    ax1.hlines(y1, 0, x1, colors=colors, linewidth=5)
+@check_figures_equal()
+def test_vlines_with_masked_colors(fig_test, fig_ref):
+    colors1 = ['red', 'green', 'blue', 'purple', 'orange', 'black']
+    fig_test, ax1 = plt.subplots()
+    x1 = np.ma.masked_equal([2, 4, 6, 8, 10, 12], 8)
+    ymin1 = np.ma.masked_equal([0, 1, -1, 0, 2, 1], 2)
+    ymax1 = np.ma.masked_equal([13, 14, 15, 16, 17, 18], 18)
+    ax1.vlines(x1, ymin1, ymax1, colors=colors1, linewidth=2)
+    ax1.set_xlim(0, 15)
+
+    colors2 = ['red', 'green', 'blue']
+    fig_ref, ax2 = plt.subplots()
+    x2 = np.asarray([2, 4, 6])
+    ymin2 = np.asarray([0, 1, -1])
+    ymax2 = np.asarray([13, 14, 15])
+    ax2.vlines(x2, ymin2, ymax2, colors=colors2, linewidth=2)
+    ax2.set_xlim(0, 15)
 
-    # GH #7406
-    y2 = [2, 3, 4, 5, 6, 7]
-    x2 = [2, -6, 3, 8, np.nan, 2]
-    fig2, (ax2, ax3, ax4) = plt.subplots(nrows=3, figsize=(4, 8))
-    ax2.hlines(y2, 0, x2, colors=colors, linewidth=5)
 
-    y3 = [2, 3, 4, 5, 6, 7]
-    x3 = [np.nan, 2, -6, 3, 8, 2]
-    ax3.hlines(y3, 0, x3, colors=colors, linewidth=3, linestyle='--')
+@check_figures_equal()
+def test_hlines_with_nan_colors(fig_test, fig_ref):
+    colors1 = ['red', 'green', 'blue', 'purple', 'orange', 'black']
+    x1 = [2, 3, 4, 5, 6, 7]
+    y1 = [2, -6, 3, 8, np.nan, 2]
 
-    y4 = [2, 3, 4, 5, 6, 7]
-    x4 = [np.nan, 2, -6, 3, 8, np.nan]
-    ax4.hlines(y4, 0, x4, colors=colors, linewidth=2)
+    fig_test, ax1 = plt.subplots()
+    ax1.vlines(y1, 0, x1, colors=colors1, linewidth=5)
 
-    # tweak the y-axis so we can see the lines better
-    for ax in [ax1, ax2, ax3, ax4]:
-        ax.set_ylim(0, 10)
+    colors2 = ['red', 'green', 'blue', 'purple', 'black']
+    x2 = [2, 3, 4, 5, 7]
+    y2 = [2, -6, 3, 8, 2]
+    # Reference image
+    fig_ref, ax2 = plt.subplots()
+    ax2.vlines(y2, 0, x2, colors=colors2, linewidth=5)
 
-    # check that the x-lims are all automatically the same
-    assert ax1.get_xlim() == ax2.get_xlim()
-    assert ax1.get_xlim() == ax3.get_xlim()
-    assert ax1.get_xlim() == ax4.get_xlim()
 
-    fig3, ax5 = plt.subplots()
-    y5 = np.ma.masked_equal([2, 4, 6, 8, 10, 12], 8)
-    xmin5 = np.ma.masked_equal([0, 1, -1, 0, 2, 1], 2)
-    xmax5 = np.ma.masked_equal([13, 14, 15, 16, 17, 18], 18)
-    ax5.hlines(y5, xmin5, xmax5, colors=colors, linewidth=2)
-    ax5.set_ylim(0, 15)
+@check_figures_equal()
+def test_hlines_with_masked_colors(fig_test, fig_ref):
+    colors1 = ['red', 'green', 'blue', 'purple', 'orange', 'black']
+    fig_test, ax1 = plt.subplots()
+    y1 = np.ma.masked_equal([2, 4, 6, 8, 10, 12], 8)
+    xmin1 = np.ma.masked_equal([0, 1, -1, 0, 2, 1], 2)
+    xmax1 = np.ma.masked_equal([13, 14, 15, 16, 17, 18], 18)
+    ax1.vlines(y1, xmin1, xmax1, colors=colors1, linewidth=2)
+    ax1.set_xlim(0, 15)
+
+    colors2 = ['red', 'green', 'blue']
+    fig_ref, ax2 = plt.subplots()
+    y2 = np.asarray([2, 4, 6])
+    xmin2 = np.asarray([0, 1, -1])
+    xmax2 = np.asarray([13, 14, 15])
+    ax2.vlines(y2, xmin2, xmax2, colors=colors2, linewidth=2)
+    ax2.set_xlim(0, 15)
 
 
 @image_comparison(['step_linestyle', 'step_linestyle'], remove_text=True)
