Merge pull request #21342 from ericpre/fix_drawing_animated_artists_callback

dstansby · web-flow · commit ea049782abe3 · 2022-01-09T10:19:36.000Z
Fix drawing animated artists changed in selector callback
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -916,6 +916,63 @@ def test_span_selector_bound(direction):
     assert tool._edge_handles.positions == handle_positions
 
 
+@pytest.mark.backend('QtAgg', skip_on_importerror=True)
+def test_span_selector_animated_artists_callback():
+    """Check that the animated artists changed in callbacks are updated."""
+    x = np.linspace(0, 2 * np.pi, 100)
+    values = np.sin(x)
+
+    fig, ax = plt.subplots()
+    (ln,) = ax.plot(x, values, animated=True)
+    (ln2, ) = ax.plot([], animated=True)
+
+    # spin the event loop to let the backend process any pending operations
+    # before drawing artists
+    # See blitting tutorial
+    plt.pause(0.1)
+    ax.draw_artist(ln)
+    fig.canvas.blit(fig.bbox)
+
+    def mean(vmin, vmax):
+        # Return mean of values in x between *vmin* and *vmax*
+        indmin, indmax = np.searchsorted(x, (vmin, vmax))
+        v = values[indmin:indmax].mean()
+        ln2.set_data(x, v)
+
+    span = widgets.SpanSelector(ax, mean, direction='horizontal',
+                                onmove_callback=mean,
+                                interactive=True,
+                                drag_from_anywhere=True,
+                                useblit=True)
+
+    # Add span selector and check that the line is draw after it was updated
+    # by the callback
+    press_data = [1, 2]
+    move_data = [2, 2]
+    do_event(span, 'press', xdata=press_data[0], ydata=press_data[1], button=1)
+    do_event(span, 'onmove', xdata=move_data[0], ydata=move_data[1], button=1)
+    assert span._get_animated_artists() == (ln, ln2)
+    assert ln.stale is False
+    assert ln2.stale
+    assert ln2.get_ydata() == 0.9547335049088455
+    span.update()
+    assert ln2.stale is False
+
+    # Change span selector and check that the line is drawn/updated after its
+    # value was updated by the callback
+    press_data = [4, 2]
+    move_data = [5, 2]
+    release_data = [5, 2]
+    do_event(span, 'press', xdata=press_data[0], ydata=press_data[1], button=1)
+    do_event(span, 'onmove', xdata=move_data[0], ydata=move_data[1], button=1)
+    assert ln.stale is False
+    assert ln2.stale
+    assert ln2.get_ydata() == -0.9424150707548072
+    do_event(span, 'release', xdata=release_data[0],
+             ydata=release_data[1], button=1)
+    assert ln2.stale is False
+
+
 def check_lasso_selector(**kwargs):
     ax = get_ax()
 
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -1839,6 +1839,18 @@ def set_active(self, active):
         if active:
             self.update_background(None)
 
+    def _get_animated_artists(self):
+        """
+        Convenience method to get all animated artists of the figure containing
+        this widget, excluding those already present in self.artists.
+        The returned tuple is not sorted by 'z_order': z_order sorting is
+        valid only when considering all artists and not only a subset of all
+        artists.
+        """
+        return tuple(a for ax_ in self.ax.get_figure().get_axes()
+                     for a in ax_.get_children()
+                     if a.get_animated() and a not in self.artists)
+
     def update_background(self, event):
         """Force an update of the background."""
         # If you add a call to `ignore` here, you'll want to check edge case:
@@ -1848,15 +1860,21 @@ def update_background(self, event):
         # Make sure that widget artists don't get accidentally included in the
         # background, by re-rendering the background if needed (and then
         # re-re-rendering the canvas with the visible widget artists).
-        needs_redraw = any(artist.get_visible() for artist in self.artists)
+        # We need to remove all artists which will be drawn when updating
+        # the selector: if we have animated artists in the figure, it is safer
+        # to redrawn by default, in case they have updated by the callback
+        # zorder needs to be respected when redrawing
+        artists = sorted(self.artists + self._get_animated_artists(),
+                         key=lambda a: a.get_zorder())
+        needs_redraw = any(artist.get_visible() for artist in artists)
         with ExitStack() as stack:
             if needs_redraw:
-                for artist in self.artists:
+                for artist in artists:
                     stack.enter_context(artist._cm_set(visible=False))
                 self.canvas.draw()
             self.background = self.canvas.copy_from_bbox(self.ax.bbox)
         if needs_redraw:
-            for artist in self.artists:
+            for artist in artists:
                 self.ax.draw_artist(artist)
 
     def connect_default_events(self):
@@ -1903,7 +1921,12 @@ def update(self):
                 self.canvas.restore_region(self.background)
             else:
                 self.update_background(None)
-            for artist in self.artists:
+            # We need to draw all artists, which are not included in the
+            # background, therefore we also draw self._get_animated_artists()
+            # and we make sure that we respect z_order
+            artists = sorted(self.artists + self._get_animated_artists(),
+                             key=lambda a: a.get_zorder())
+            for artist in artists:
                 self.ax.draw_artist(artist)
             self.canvas.blit(self.ax.bbox)
         else:
