Improve draw performance

ultraplot/_animation.py

@@ -0,0 +1,231 @@
+#!/usr/bin/env python3
+"""
+Helpers for responsive interactive and animated UltraPlot figures.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from contextlib import contextmanager
+from weakref import WeakSet
+
+import matplotlib.artist as martist
+import matplotlib.transforms as mtransforms
+
+
+class _BlitManager:
+    """
+    Manage efficient updates of a small set of changing artists.
+
+    The manager caches the static canvas background, restores it for each
+    update, redraws only the managed artists, and blits the affected region.
+    Backends without blitting support safely fall back to ``draw_idle()``.
+
+    Parameters
+    ----------
+    canvas : `~matplotlib.backend_bases.FigureCanvasBase`
+        Canvas containing the artists.
+    artists : iterable of `~matplotlib.artist.Artist`, optional
+        Artists that will change between updates.
+    bbox : `~matplotlib.transforms.Bbox` or object with a ``bbox`` attribute, optional
+        Region to cache and blit. By default, the union of the managed artists'
+        axes bounding boxes is used. Figure-level artists fall back to the full
+        figure bounding box.
+
+    Notes
+    -----
+    Managed artists are drawn above the cached static background, matching
+    Matplotlib's standard blitting behavior.
+    """
+
+    def __init__(self, canvas, artists: Iterable[martist.Artist] = (), bbox=None):
+        if canvas.figure is None:
+            raise RuntimeError("Cannot manage blitting for a canvas without a figure.")
+        self.canvas = canvas
+        self.figure = canvas.figure
+        self._artists = []
+        self._animated = {}
+        self._bbox = bbox
+        self._background = None
+        self._closed = False
+        self._suspended = False
+        self._supports_blit = bool(canvas.supports_blit)
+        self._draw_cid = canvas.mpl_connect("draw_event", self._on_draw)
+        self._resize_cid = canvas.mpl_connect("resize_event", self._on_resize)
+        managers = getattr(self.figure, "_blit_managers", None)
+        if managers is None:
+            managers = self.figure._blit_managers = WeakSet()
+        managers.add(self)
+        for artist in artists:
+            self.add_artist(artist)
+
+    @property
+    def artists(self):
+        """Managed artists as an immutable tuple."""
+        return tuple(self._artists)
+
+    @property
+    def supports_blit(self):
+        """Whether the associated canvas supports blitting."""
+        return self._supports_blit
+
+    def _resolve_bbox(self):
+        bbox = self._bbox
+        if bbox is not None:
+            return getattr(bbox, "bbox", bbox)
+
+        axes = []
+        for artist in self._artists:
+            axis = getattr(artist, "axes", None)
+            if axis is None:
+                return self.figure.bbox
+            if axis not in axes:
+                axes.append(axis)
+        if not axes:
+            return self.figure.bbox
+        return mtransforms.Bbox.union([axis.bbox for axis in axes])
+
+    def _draw_artists(self):
+        for artist in sorted(self._artists, key=lambda item: item.get_zorder()):
+            self.figure.draw_artist(artist)
+
+    def _on_draw(self, event):
+        if self._closed or self._suspended or not self._supports_blit:
+            return
+        if event is not None and event.canvas is not self.canvas:
+            return
+        self._background = self.canvas.copy_from_bbox(self._resolve_bbox())
+        self._draw_artists()
+
+    def _on_resize(self, event):
+        if event is None or event.canvas is self.canvas:
+            self.invalidate()
+
+    def add_artist(self, artist: martist.Artist):
+        """
+        Add an artist to the managed update set.
+
+        Returns
+        -------
+        _BlitManager
+            This manager, to permit chained calls.
+        """
+        if self._closed:
+            raise RuntimeError("Cannot add artists to a closed _BlitManager.")
+        if not isinstance(artist, martist.Artist):
+            raise TypeError(
+                f"Expected a matplotlib Artist, got {type(artist).__name__}."
+            )
+        if artist.figure is not self.figure:
+            raise RuntimeError("The artist must belong to the manager's figure.")
+        if artist in self._artists:
+            return self
+        self._artists.append(artist)
+        self._animated[artist] = artist.get_animated()
+        if self._supports_blit:
+            artist.set_animated(True)
+        self.invalidate()
+        return self
+
+    def remove_artist(self, artist: martist.Artist):
+        """
+        Stop managing an artist and restore its original animated state.
+
+        Returns
+        -------
+        _BlitManager
+            This manager, to permit chained calls.
+        """
+        if artist not in self._artists:
+            return self
+        self._artists.remove(artist)
+        artist.set_animated(self._animated.pop(artist))
+        self.invalidate()
+        return self
+
+    def invalidate(self):
+        """Discard the cached background before the next update."""
+        self._background = None
+
+    @contextmanager
+    def _save_context(self):
+        """Temporarily restore original artist states for a complete export."""
+        if self._closed:
+            yield
+            return
+        self._suspended = True
+        current = {artist: artist.get_animated() for artist in self._artists}
+        try:
+            for artist in self._artists:
+                artist.set_animated(self._animated[artist])
+            yield
+        finally:
+            for artist, animated in current.items():
+                artist.set_animated(animated)
+            self._suspended = False
+            self.invalidate()
+
+    def update(self, *, flush=False):
+        """
+        Redraw the managed artists.
+
+        Parameters
+        ----------
+        flush : bool, default: False
+            Whether to immediately process pending GUI events after blitting.
+
+        Returns
+        -------
+        bool
+            ``True`` when the blitting fast path was used, otherwise ``False``.
+        """
+        if self._closed:
+            raise RuntimeError("Cannot update a closed _BlitManager.")
+        if not self._supports_blit:
+            self.canvas.draw_idle()
+            if flush:
+                self.canvas.flush_events()
+            return False
+
+        if self._background is None:
+            # The draw_event callback captures the static background and draws
+            # the animated artists before the backend presents the frame.
+            self.canvas.draw()
+        else:
+            bbox = self._resolve_bbox()
+            self.canvas.restore_region(self._background)
+            self._draw_artists()
+            self.canvas.blit(bbox)
+        if flush:
+            self.canvas.flush_events()
+        return True
+
+    def close(self, *, redraw=True):
+        """
+        Disconnect callbacks and restore the artists' animated states.
+
+        Parameters
+        ----------
+        redraw : bool, default: True
+            Whether to schedule a normal full redraw after restoring the artists.
+        """
+        if self._closed:
+            return
+        self.canvas.mpl_disconnect(self._draw_cid)
+        self.canvas.mpl_disconnect(self._resize_cid)
+        for artist in tuple(self._artists):
+            artist.set_animated(self._animated[artist])
+        self._artists.clear()
+        self._animated.clear()
+        self._background = None
+        self._closed = True
+        managers = getattr(self.figure, "_blit_managers", ())
+        managers.discard(self)
+        if redraw:
+            self.canvas.draw_idle()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, traceback):
+        self.close()