BUG: fix .remove method for container when one of the items is None, which occurs for errorbars when fmt='none'

astrofrog · astrofrog · commit 994076f78f26 · 2017-08-09T20:46:37.000+01:00
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -655,7 +655,7 @@ def flatten(seq, scalarp=is_scalar_or_string):
     and Recipe 1.12 in cookbook
     """
     for item in seq:
-        if scalarp(item):
+        if scalarp(item) or item is None:
             yield item
         else:
             for subitem in flatten(item, scalarp):
diff --git a/lib/matplotlib/container.py b/lib/matplotlib/container.py
@@ -34,7 +34,8 @@ def set_remove_method(self, f):
     def remove(self):
         for c in cbook.flatten(
                 self, scalarp=lambda x: isinstance(x, martist.Artist)):
-            c.remove()
+            if c is not None:
+                c.remove()
 
         if self._remove_method:
             self._remove_method(self)
diff --git a/lib/matplotlib/tests/test_container.py b/lib/matplotlib/tests/test_container.py
@@ -9,3 +9,26 @@ def test_stem_remove():
     ax = plt.gca()
     st = ax.stem([1, 2], [1, 2])
     st.remove()
+
+
+def test_errorbar_remove():
+
+    # Regression test for a bug that caused remove to fail when using
+    # fmt='none'
+
+    ax = plt.gca()
+
+    eb = ax.errorbar([1], [1])
+    eb.remove()
+
+    eb = ax.errorbar([1], [1], xerr=1)
+    eb.remove()
+
+    eb = ax.errorbar([1], [1], yerr=2)
+    eb.remove()
+
+    eb = ax.errorbar([1], [1], xerr=[2], yerr=2)
+    eb.remove()
+
+    eb = ax.errorbar([1], [1], fmt='none')
+    eb.remove()
