Merge pull request #8737 from dstansby/fix-tests

anntzer · web-flow · commit 166a14473272 · 2017-06-12T10:16:29.000-07:00
Fix colorbar test and color level determination for contour
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -1254,11 +1254,11 @@ def _process_colors(self):
             i0, i1 = 0, len(self.levels)
             if self.filled:
                 i1 -= 1
-            # Out of range indices for over and under:
-            if self.extend in ('both', 'min'):
-                i0 = -1
-            if self.extend in ('both', 'max'):
-                i1 += 1
+                # Out of range indices for over and under:
+                if self.extend in ('both', 'min'):
+                    i0 -= 1
+                if self.extend in ('both', 'max'):
+                    i1 += 1
             self.cvalues = list(range(i0, i1))
             self.set_norm(colors.NoNorm())
         else:
diff --git a/lib/matplotlib/tests/baseline_images/test_contour/contour_manual_colors_and_levels.png b/lib/matplotlib/tests/baseline_images/test_contour/contour_manual_colors_and_levels.png
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -168,21 +168,22 @@ def test_given_colors_levels_and_extends():
     levels = [2, 4, 8, 10]
 
     for i, ax in enumerate(axes.flatten()):
-        plt.sca(ax)
-
         filled = i % 2 == 0.
         extend = ['neither', 'min', 'max', 'both'][i // 2]
 
         if filled:
-            last_color = -1 if extend in ['min', 'max'] else None
-            plt.contourf(data, colors=colors[:last_color], levels=levels,
-                         extend=extend)
+            # If filled, we have 3 colors with no extension,
+            # 4 colors with one extension, and 5 colors with both extensions
+            first_color = 1 if extend in ['max', 'neither'] else None
+            last_color = -1 if extend in ['min', 'neither'] else None
+            c = ax.contourf(data, colors=colors[first_color:last_color],
+                            levels=levels, extend=extend)
         else:
-            last_level = -1 if extend == 'both' else None
-            plt.contour(data, colors=colors, levels=levels[:last_level],
-                        extend=extend)
+            # If not filled, we have 4 levels and 4 colors
+            c = ax.contour(data, colors=colors[:-1],
+                           levels=levels, extend=extend)
 
-        plt.colorbar()
+        plt.colorbar(c, ax=ax)
 
 
 @image_comparison(baseline_images=['contour_datetime_axis'],
