Merge consecutive rasterizations

samtygier · samtygier · commit b6a273989ffc · 2020-05-03T17:15:45.000+01:00
In vector output it is possible to flag artists to be rasterized. In many cases with multiple rasterized objects there can be significant file size savings by combining the rendered bitmaps into a single bitmap. This is achieved by moving the depth tracking logic from start_rasterizing() and stop_rasterizing() functions into the allow_rasterization() wrapper. This allows delaying the call to stop_rasterizing() until we are about to draw an non rasterized artist. The outer draw method, i.e. in Figure must be wraped with _finalize_rasterization() to ensure the that rasterization is completed. Figure.suppressComposite can be used to prevent merging. This fixes #17149
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -34,7 +34,17 @@ def allow_rasterization(draw):
     def draw_wrapper(artist, renderer, *args, **kwargs):
         try:
             if artist.get_rasterized():
-                renderer.start_rasterizing()
+                if renderer._raster_depth == 0 and not renderer._rasterizing:
+                    renderer.start_rasterizing()
+                    renderer._rasterizing = True
+                renderer._raster_depth += 1
+            else:
+                if renderer._raster_depth == 0 and renderer._rasterizing:
+                    # Only stop when we are not in a rasterized parent
+                    # and something has be rasterized since last stop
+                    renderer.stop_rasterizing()
+                    renderer._rasterizing = False
+
             if artist.get_agg_filter() is not None:
                 renderer.start_filter()
 
@@ -43,12 +53,32 @@ def draw_wrapper(artist, renderer, *args, **kwargs):
             if artist.get_agg_filter() is not None:
                 renderer.stop_filter(artist.get_agg_filter())
             if artist.get_rasterized():
+                renderer._raster_depth -= 1
+            if (renderer._rasterizing and artist.figure and
+                    artist.figure.suppressComposite):
+                # restart rasterizing to prevent merging
                 renderer.stop_rasterizing()
+                renderer.start_rasterizing()
 
     draw_wrapper._supports_rasterization = True
     return draw_wrapper
 
 
+def _finalize_rasterization(draw):
+    """
+    Decorator for Artist.draw method. Needed on the outermost artist, i.e.
+    Figure, to finish up if the render is still in rasterized mode.
+    """
+    @wraps(draw)
+    def draw_wrapper(artist, renderer, *args, **kwargs):
+        result = draw(artist, renderer, *args, **kwargs)
+        if renderer._rasterizing:
+            renderer.stop_rasterizing()
+            renderer._rasterizing = False
+        return result
+    return draw_wrapper
+
+
 def _stale_axes_callback(self, val):
     if self.axes:
         self.axes.stale = val
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -142,6 +142,8 @@ def __init__(self):
         super().__init__()
         self._texmanager = None
         self._text2path = textpath.TextToPath()
+        self._raster_depth = 0
+        self._rasterizing = False
 
     def open_group(self, s, gid=None):
         """
diff --git a/lib/matplotlib/backends/backend_mixed.py b/lib/matplotlib/backends/backend_mixed.py
@@ -52,7 +52,6 @@ def __init__(self, figure, width, height, dpi, vector_renderer,
         self._vector_renderer = vector_renderer
 
         self._raster_renderer = None
-        self._rasterizing = 0
 
         # A reference to the figure is needed as we need to change
         # the figure dpi before and after the rasterization. Although
@@ -84,47 +83,40 @@ def start_rasterizing(self):
             r = process_figure_for_rasterizing(self.figure,
                                                self._bbox_inches_restore)
             self._bbox_inches_restore = r
-        if self._rasterizing == 0:
-            self._raster_renderer = self._raster_renderer_class(
-                self._width*self.dpi, self._height*self.dpi, self.dpi)
-            self._renderer = self._raster_renderer
-        self._rasterizing += 1
+
+        self._raster_renderer = self._raster_renderer_class(
+            self._width*self.dpi, self._height*self.dpi, self.dpi)
+        self._renderer = self._raster_renderer
 
     def stop_rasterizing(self):
         """
         Exit "raster" mode.  All of the drawing that was done since
         the last `start_rasterizing` call will be copied to the
         vector backend by calling draw_image.
-
-        If `start_rasterizing` has been called multiple times,
-        `stop_rasterizing` must be called the same number of times before
-        "raster" mode is exited.
         """
-        self._rasterizing -= 1
-        if self._rasterizing == 0:
-            self._renderer = self._vector_renderer
-
-            height = self._height * self.dpi
-            buffer, bounds = self._raster_renderer.tostring_rgba_minimized()
-            l, b, w, h = bounds
-            if w > 0 and h > 0:
-                image = np.frombuffer(buffer, dtype=np.uint8)
-                image = image.reshape((h, w, 4))
-                image = image[::-1]
-                gc = self._renderer.new_gc()
-                # TODO: If the mixedmode resolution differs from the figure's
-                #       dpi, the image must be scaled (dpi->_figdpi). Not all
-                #       backends support this.
-                self._renderer.draw_image(
-                    gc,
-                    l * self._figdpi / self.dpi,
-                    (height-b-h) * self._figdpi / self.dpi,
-                    image)
-            self._raster_renderer = None
-            self._rasterizing = False
-
-            # restore the figure dpi.
-            self.figure.set_dpi(self._figdpi)
+
+        self._renderer = self._vector_renderer
+
+        height = self._height * self.dpi
+        buffer, bounds = self._raster_renderer.tostring_rgba_minimized()
+        l, b, w, h = bounds
+        if w > 0 and h > 0:
+            image = np.frombuffer(buffer, dtype=np.uint8)
+            image = image.reshape((h, w, 4))
+            image = image[::-1]
+            gc = self._renderer.new_gc()
+            # TODO: If the mixedmode resolution differs from the figure's
+            #       dpi, the image must be scaled (dpi->_figdpi). Not all
+            #       backends support this.
+            self._renderer.draw_image(
+                gc,
+                l * self._figdpi / self.dpi,
+                (height-b-h) * self._figdpi / self.dpi,
+                image)
+        self._raster_renderer = None
+
+        # restore the figure dpi.
+        self.figure.set_dpi(self._figdpi)
 
         if self._bbox_inches_restore:  # when tight bbox is used
             r = process_figure_for_rasterizing(self.figure,
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -18,7 +18,8 @@
 from matplotlib import __version__ as _mpl_version
 
 import matplotlib.artist as martist
-from matplotlib.artist import Artist, allow_rasterization
+from matplotlib.artist import (
+    Artist, allow_rasterization, _finalize_rasterization)
 from matplotlib.backend_bases import (
     FigureCanvasBase, NonGuiException, MouseButton)
 import matplotlib.cbook as cbook
@@ -1689,6 +1690,7 @@ def clear(self, keep_observers=False):
         """Clear the figure -- synonym for `clf`."""
         self.clf(keep_observers=keep_observers)
 
+    @_finalize_rasterization
     @allow_rasterization
     def draw(self, renderer):
         # docstring inherited
