Merge pull request #1229 from ajdawson/colorbar-triangles

NF - option to make colorbar extensions rectangles

doc/api/api_changes.rst

@@ -11,6 +11,14 @@ help figure out possible sources of the changes you are experiencing.
 For new features that were added to matplotlib, please see
 :ref:`whats-new`.
 
+
+Changes in 1.3.x
+================
+
+* A new keyword *extendrect* in :meth:`~matplotlib.pyplot.colorbar` and
+  :class:`~matplotlib.colorbar.ColorbarBase` allows one to control the shape
+  of colorbar extensions.
+
 Changes in 1.2.x
 ================
 

doc/users/whats_new.rst

@@ -15,6 +15,18 @@ revision, see the :ref:`github-stats`.
    versions 2.4 to 2.7.  matplotlib 1.2 and later require
    versions 2.6, 2.7, and 3.1 and higher.
 
+.. _whats-new-1-3:
+
+new in matplotlib-1.3
+=====================
+
+Rectangular colorbar extensions
+-------------------------------
+Andrew Dawson added a new keyword argument *extendrect* to
+:meth:`~matplotlib.pyplot.colorbar` to optionally make colorbar
+extensions rectangular instead of triangular.
+
+
 .. _whats-new-1-2:
 
 new in matplotlib-1.2

lib/matplotlib/colorbar.py

@@ -82,6 +82,10 @@
                   be given, indicating the lengths of the minimum and
                   maximum colorbar extensions respectively as a
                   fraction of the interior colorbar length.
+    *extendrect*  [ *False* | *True* ]
+                  If *False* the minimum and maximum colorbar extensions
+                  will be triangular (the default). If *True* the
+                  extensions will be rectangular.
     *spacing*     [ 'uniform' | 'proportional' ]
                   Uniform spacing gives each discrete color the same
                   space; proportional makes the space proportional to
@@ -258,6 +262,7 @@ def __init__(self, ax, cmap=None,
                  drawedges=False,
                  filled=True,
                  extendfrac=None,
+                 extendrect=False,
                  ):
         self.ax = ax
         self._patch_ax()
@@ -276,6 +281,7 @@ def __init__(self, ax, cmap=None,
         self.drawedges = drawedges
         self.filled = filled
         self.extendfrac = extendfrac
+        self.extendrect = extendrect
         self.solids = None
         self.lines = list()
         self.outline = None
@@ -773,9 +779,9 @@ def _mesh(self):
             y = self._proportional_y()
         self._y = y
         X, Y = np.meshgrid(x, y)
-        if self._extend_lower():
+        if self._extend_lower() and not self.extendrect:
             X[0, :] = 0.5
-        if self._extend_upper():
+        if self._extend_upper() and not self.extendrect:
             X[-1, :] = 0.5
         return X, Y
 

lib/matplotlib/tests/test_colorbar.py

@@ -6,23 +6,69 @@
 from matplotlib.colorbar import ColorbarBase
 
 
-def _colorbar_extensions(spacing):
+def _get_cmap_norms():
+    """
+    Define a colormap and appropriate norms for each of the four
+    possible settings of the extend keyword.
 
+    Helper function for _colorbar_extension_shape and
+    colorbar_extension_length.
+    """
     # Create a color map and specify the levels it represents.
     cmap = get_cmap("RdBu", lut=5)
     clevs = [-5., -2.5, -.5, .5, 1.5, 3.5]
-
     # Define norms for the color maps.
     norms = dict()
     norms['neither'] = BoundaryNorm(clevs, len(clevs)-1)
     norms['min'] = BoundaryNorm([-10]+clevs[1:], len(clevs)-1)
     norms['max'] = BoundaryNorm(clevs[:-1]+[10], len(clevs)-1)
     norms['both'] = BoundaryNorm([-10]+clevs[1:-1]+[10], len(clevs)-1)
+    return cmap, norms
+
+
+def _colorbar_extension_shape(spacing):
+    '''
+    Produce 4 colorbars with rectangular extensions for either uniform
+    or proportional spacing.
 
+    Helper function for test_colorbar_extension_shape.
+    '''
+    # Get a colormap and appropriate norms for each extension type.
+    cmap, norms = _get_cmap_norms()
+    # Create a figure and adjust whitespace for subplots.
+    fig = plt.figure()
+    fig.subplots_adjust(hspace=4)
+    for i, extension_type in enumerate(('neither', 'min', 'max', 'both')):
+        # Get the appropriate norm and use it to get colorbar boundaries.
+        norm = norms[extension_type]
+        boundaries = values = norm.boundaries
+        # Create a subplot.
+        cax = fig.add_subplot(4, 1, i+1)
+        # Turn off text and ticks.
+        for item in cax.get_xticklabels() + cax.get_yticklabels() +\
+                cax.get_xticklines() + cax.get_yticklines():
+            item.set_visible(False)
+        # Generate the colorbar.
+        cb = ColorbarBase(cax, cmap=cmap, norm=norm,
+                boundaries=boundaries, values=values,
+                extend=extension_type, extendrect=True,
+                orientation='horizontal', spacing=spacing)
+    # Return the figure to the caller.
+    return fig
+
+
+def _colorbar_extension_length(spacing):
+    '''
+    Produce 12 colorbars with variable length extensions for either
+    uniform or proportional spacing.
+
+    Helper function for test_colorbar_extension_length.
+    '''
+    # Get a colormap and appropriate norms for each extension type.
+    cmap, norms = _get_cmap_norms()
     # Create a figure and adjust whitespace for subplots.
     fig = plt.figure()
     fig.subplots_adjust(hspace=.6)
-
     for i, extension_type in enumerate(('neither', 'min', 'max', 'both')):
         # Get the appropriate norm and use it to get colorbar boundaries.
         norm = norms[extension_type]
@@ -39,20 +85,33 @@ def _colorbar_extensions(spacing):
                     boundaries=boundaries, values=values,
                     extend=extension_type, extendfrac=extendfrac,
                     orientation='horizontal', spacing=spacing)
-
     # Return the figure to the caller.
     return fig
 
 
+@image_comparison(
+        baseline_images=['colorbar_extensions_shape_uniform',
+                         'colorbar_extensions_shape_proportional'],
+        extensions=['png'])
+def test_colorbar_extension_shape():
+    '''Test rectangular colorbar extensions.'''
+    # Use default params so matplotlibrc doesn't cause the test to fail.
+    rcParams.update(rcParamsDefault)
+    # Create figures for uniform and proportionally spaced colorbars.
+    fig1 = _colorbar_extension_shape('uniform')
+    fig2 = _colorbar_extension_shape('proportional')
+
+
 @image_comparison(
         baseline_images=['colorbar_extensions_uniform', 'colorbar_extensions_proportional'],
         extensions=['png'])
-def test_colorbar_extensions():
-    # Use default params so .matplotlibrc doesn't cause the test to fail.
+def test_colorbar_extension_length():
+    '''Test variable length colorbar extensions.'''
+    # Use default params so matplotlibrc doesn't cause the test to fail.
     rcParams.update(rcParamsDefault)
     # Create figures for uniform and proportionally spaced colorbars.
-    fig1 = _colorbar_extensions('uniform')
-    fig2 = _colorbar_extensions('proportional')
+    fig1 = _colorbar_extension_length('uniform')
+    fig2 = _colorbar_extension_length('proportional')
 
 
 if __name__ == '__main__':