Adding Colorbar.remove()

Author: KevKeating

CHANGELOG

@@ -1,3 +1,5 @@
+2014-06-10 Added Colorbar.remove()
+
 2014-06-07 Fixed bug so radial plots can be saved as ps in py3k.
 
 2014-06-01 Changed the fmt kwarg of errorbar to support the

lib/matplotlib/colorbar.py

@@ -840,6 +840,14 @@ def _locate(self, x):
     def set_alpha(self, alpha):
         self.alpha = alpha
 
+    def remove(self):
+        """
+        Remove this colorbar from the figure
+        """
+
+        fig = self.ax.figure
+        fig.delaxes(self.ax)
+
 
 class Colorbar(ColorbarBase):
     """
@@ -967,6 +975,33 @@ def update_bruteforce(self, mappable):
         # be recalculating everything if there was a simple alpha
         # change.
 
+    def remove(self):
+        """
+        Remove this colorbar from the figure.  If the colorbar was created with
+        ``use_gridspec=True`` then restore the gridspec to its previous value.
+        """
+
+        ColorbarBase.remove(self)
+        self.mappable.callbacksSM.disconnect(self.mappable.colorbar_cid)
+        self.mappable.colorbar = None
+        self.mappable.colorbar_cid = None
+
+        try:
+            ax = self.mappable.axes
+        except AttributeError:
+            return
+
+        try:
+            gs = ax.get_subplotspec().get_gridspec()
+            subplotspec = gs.get_topmost_subplotspec()
+        except AttributeError:
+            # use_gridspec was False
+            pos = ax.get_position(original=True)
+            ax.set_position(pos)
+        else:
+            # use_gridspec was True
+            ax.set_subplotspec(subplotspec)
+
 
 @docstring.Substitution(make_axes_kw_doc)
 def make_axes(parents, location=None, orientation=None, fraction=0.15,
@@ -1280,7 +1315,8 @@ def colorbar_factory(cax, mappable, **kwargs):
     else:
         cb = Colorbar(cax, mappable, **kwargs)
 
-    mappable.callbacksSM.connect('changed', cb.on_mappable_changed)
+    cid = mappable.callbacksSM.connect('changed', cb.on_mappable_changed)
     mappable.colorbar = cb
+    mappable.colorbar_cid = cid
 
     return cb

lib/matplotlib/tests/test_colorbar.py

@@ -6,7 +6,7 @@
 import numpy as np
 from numpy import ma
 import matplotlib
-from matplotlib.testing.decorators import image_comparison, knownfailureif
+from matplotlib.testing.decorators import image_comparison, cleanup
 import matplotlib.pyplot as plt
 from matplotlib import rcParams
 from matplotlib.colors import BoundaryNorm
@@ -211,6 +211,41 @@ def test_colorbar_single_scatter():
     plt.colorbar(cs)
 
 
+def _test_remove_from_figure(use_gridspec):
+    """
+    Test `remove_from_figure` with the specified ``use_gridspec`` setting
+    """
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    sc = ax.scatter([1, 2], [3, 4], cmap="spring")
+    sc.set_array(np.array([5, 6]))
+    pre_figbox = np.array(ax.figbox)
+    cb = fig.colorbar(sc, use_gridspec=use_gridspec)
+    fig.subplots_adjust()
+    cb.remove()
+    fig.subplots_adjust()
+    post_figbox = np.array(ax.figbox)
+    assert (pre_figbox == post_figbox).all()
+
+
+@cleanup
+def test_remove_from_figure_with_gridspec():
+    """
+    Make sure that `remove_from_figure` removes the colorbar and properly
+    restores the gridspec
+    """
+    _test_remove_from_figure(True)
+
+
+@cleanup
+def test_remove_from_figure_no_gridspec():
+    """
+    Make sure that `remove_from_figure` removes a colorbar that was created
+    without modifying the gridspec
+    """
+    _test_remove_from_figure(False)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)