Merge pull request #4969 from tacaswell/fix_line2D_remove

WeatherGod · WeatherGod · commit 09e044cf0517 · 2015-08-19T23:00:39.000-04:00
FIX: account for None in Line2D.axes setter
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -568,13 +568,14 @@ def get_window_extent(self, renderer):
     def axes(self, ax):
         # call the set method from the base-class property
         Artist.axes.fset(self, ax)
-        # connect unit-related callbacks
-        if ax.xaxis is not None:
-            self._xcid = ax.xaxis.callbacks.connect('units',
-                                                    self.recache_always)
-        if ax.yaxis is not None:
-            self._ycid = ax.yaxis.callbacks.connect('units',
-                                                    self.recache_always)
+        if ax is not None:
+            # connect unit-related callbacks
+            if ax.xaxis is not None:
+                self._xcid = ax.xaxis.callbacks.connect('units',
+                                                        self.recache_always)
+            if ax.yaxis is not None:
+                self._ycid = ax.yaxis.callbacks.connect('units',
+                                                        self.recache_always)
 
     def set_data(self, *args):
         """
diff --git a/lib/matplotlib/tests/test_artist.py b/lib/matplotlib/tests/test_artist.py
@@ -150,21 +150,27 @@ def test_cull_markers():
 def test_remove():
     fig, ax = plt.subplots()
     im = ax.imshow(np.arange(36).reshape(6, 6))
+    ln, = ax.plot(range(5))
 
     assert_true(fig.stale)
     assert_true(ax.stale)
 
     fig.canvas.draw()
     assert_false(fig.stale)
     assert_false(ax.stale)
+    assert_false(ln.stale)
 
     assert_true(im in ax.mouseover_set)
+    assert_true(ln not in ax.mouseover_set)
     assert_true(im.axes is ax)
 
     im.remove()
+    ln.remove()
+
+    for art in [im, ln]:
+        assert_true(art.axes is None)
+        assert_true(art.figure is None)
 
-    assert_true(im.axes is None)
-    assert_true(im.figure is None)
     assert_true(im not in ax.mouseover_set)
     assert_true(fig.stale)
     assert_true(ax.stale)
