Refactor image compositing logic into a single function

mdboom · mdboom · commit b391f0844342 · 2016-02-17T09:47:57.000-05:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2352,38 +2352,7 @@ def draw(self, renderer=None, inframe=False):
                 a.draw(renderer)
             renderer.stop_rasterizing()
 
-        has_images = np.any(isinstance(x[1], mimage._ImageBase) for x in dsu)
-
-        if renderer.option_image_nocomposite() or not has_images:
-            for zorder, a in dsu:
-                a.draw(renderer)
-        else:
-            # Composite any adjacent images together
-            image_group = []
-            mag = renderer.get_image_magnification()
-
-            def flush_images():
-                if len(image_group) == 1:
-                    image_group[0].draw(renderer)
-                elif len(image_group) > 1:
-                    data, l, b = mimage.composite_images(
-                        image_group, renderer, mag)
-                    gc = renderer.new_gc()
-                    gc.set_clip_rectangle(self.bbox)
-                    gc.set_clip_path(mtransforms.TransformedPath(
-                        self.patch.get_path(),
-                        self.patch.get_transform()))
-                    renderer.draw_image(gc, round(l), round(b), data)
-                    gc.restore()
-                del image_group[:]
-
-            for zorder, a in dsu:
-                if isinstance(a, mimage._ImageBase) and a.can_composite():
-                    image_group.append(a)
-                else:
-                    flush_images()
-                    a.draw(renderer)
-            flush_images()
+        mimage._draw_list_compositing_images(renderer, self, dsu)
 
         renderer.close_group('axes')
         self._cachedRenderer = renderer
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -1233,65 +1233,33 @@ def draw(self, renderer):
         dsu = []
 
         for a in self.patches:
-            dsu.append((a.get_zorder(), a, a.draw, [renderer]))
+            dsu.append((a.get_zorder(), a))
 
         for a in self.lines:
-            dsu.append((a.get_zorder(), a, a.draw, [renderer]))
+            dsu.append((a.get_zorder(), a))
 
         for a in self.artists:
-            dsu.append((a.get_zorder(), a, a.draw, [renderer]))
-
-        # override the renderer default if self.suppressComposite
-        # is not None
-        not_composite = renderer.option_image_nocomposite()
-        if self.suppressComposite is not None:
-            not_composite = self.suppressComposite
+            dsu.append((a.get_zorder(), a))
 
         for a in self.images:
-            dsu.append((a.get_zorder(), a, a.draw, [renderer]))
+            dsu.append((a.get_zorder(), a))
 
         # render the axes
         for a in self.axes:
-            dsu.append((a.get_zorder(), a, a.draw, [renderer]))
+            dsu.append((a.get_zorder(), a))
 
         # render the figure text
         for a in self.texts:
-            dsu.append((a.get_zorder(), a, a.draw, [renderer]))
+            dsu.append((a.get_zorder(), a))
 
         for a in self.legends:
-            dsu.append((a.get_zorder(), a, a.draw, [renderer]))
+            dsu.append((a.get_zorder(), a))
 
         dsu = [row for row in dsu if not row[1].get_animated()]
         dsu.sort(key=itemgetter(0))
 
-        if not_composite or len(self.images) == 0:
-            for zorder, a, func, args in dsu:
-                func(*args)
-        else:
-            # Composite any adjacent images together
-            image_group = []
-            mag = renderer.get_image_magnification()
-
-            def flush_images():
-                if len(image_group) == 1:
-                    image_group[0].draw(renderer)
-                elif len(image_group) > 1:
-                    data, l, b = mimage.composite_images(
-                        image_group, renderer, mag)
-                    gc = renderer.new_gc()
-                    gc.set_clip_rectangle(self.bbox)
-                    gc.set_clip_path(self.get_clip_path())
-                    renderer.draw_image(gc, l, b, data)
-                    gc.restore()
-                del image_group[:]
-
-            for zorder, a, func, args in dsu:
-                if isinstance(a, mimage._ImageBase) and a.can_composite():
-                    image_group.append(a)
-                else:
-                    flush_images()
-                    func(*args)
-            flush_images()
+        mimage._draw_list_compositing_images(
+            renderer, self, dsu, self.suppressComposite)
 
         renderer.close_group('figure')
         self.stale = False
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -117,6 +117,55 @@ def composite_images(images, renderer, magnification=1.0):
     return output, bbox.x0 / magnification, bbox.y0 / magnification
 
 
+def _draw_list_compositing_images(
+        renderer, parent, dsu, suppress_composite=None):
+    """
+    Draw a sorted list of artists, compositing images into a single
+    image where possible.
+
+    For internal matplotlib use only: It is here to reduce duplication
+    between `Figure.draw` and `Axes.draw`, but otherwise should not be
+    generally useful.
+    """
+    has_images = np.any(isinstance(x[1], _ImageBase) for x in dsu)
+
+    # override the renderer default if self.suppressComposite
+    # is not None
+    not_composite = renderer.option_image_nocomposite()
+    if suppress_composite is not None:
+        not_composite = suppress_composite
+
+    if not_composite or not has_images:
+        for zorder, a in dsu:
+            a.draw(renderer)
+    else:
+        # Composite any adjacent images together
+        image_group = []
+        mag = renderer.get_image_magnification()
+
+        def flush_images():
+            if len(image_group) == 1:
+                image_group[0].draw(renderer)
+            elif len(image_group) > 1:
+                data, l, b = composite_images(
+                    image_group, renderer, mag)
+                gc = renderer.new_gc()
+                gc.set_clip_rectangle(parent.bbox)
+                gc.set_clip_path(parent.get_clip_path())
+                renderer.draw_image(gc, round(l), round(b), data)
+                gc.restore()
+            del image_group[:]
+
+        for zorder, a in dsu:
+            if isinstance(a, _ImageBase) and a.can_composite():
+                image_group.append(a)
+            else:
+                flush_images()
+                a.draw(renderer)
+        flush_images()
+
+
+
 def _rgb_to_rgba(A):
     """
     Convert an RGB image to RGBA, as required by the image resample C++
