matplotlib · WeatherGod · Aug 20, 2015 · Aug 20, 2015
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)