FIX: Container remove method

tacaswell · tacaswell · commit 4604517b1615 · 2015-12-21T21:23:35.000-05:00
It was not properly looping over sub-children of the container. closes #5692
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1440,6 +1440,7 @@ def tk_window_focus():
     'matplotlib.tests.test_colorbar',
     'matplotlib.tests.test_colors',
     'matplotlib.tests.test_compare_images',
+    'matplotlib.tests.test_container',
     'matplotlib.tests.test_contour',
     'matplotlib.tests.test_dates',
     'matplotlib.tests.test_delaunay',
diff --git a/lib/matplotlib/container.py b/lib/matplotlib/container.py
@@ -4,6 +4,7 @@
 from matplotlib.externals import six
 
 import matplotlib.cbook as cbook
+import matplotlib.artist as martist
 
 
 class Container(tuple):
@@ -31,7 +32,9 @@ def set_remove_method(self, f):
         self._remove_method = f
 
     def remove(self):
-        for c in self:
+        for c in cbook.flatten(self,
+                               scalarp=lambda x: isinstance(x,
+                                                            martist.Artist)):
             c.remove()
 
         if self._remove_method:
diff --git a/lib/matplotlib/tests/test_container.py b/lib/matplotlib/tests/test_container.py
@@ -0,0 +1,14 @@
+from __future__ import (absolute_import, division, print_function,
+                        unicode_literals)
+
+from matplotlib.externals import six
+import matplotlib.pyplot as plt
+
+from matplotlib.testing.decorators import cleanup
+
+
+@cleanup
+def test_stem_remove():
+    ax = plt.gca()
+    st = ax.stem([1, 2], [1, 2])
+    st.remove()
