Convert adjust_bbox to use ExitStack.

QuLogic · QuLogic · commit 602be075297a · 2020-06-10T00:46:02.000-04:00
This allows us to use `_setattr_cm` and not worry about manually
restoring things or handling weird attributes ourselves.
diff --git a/lib/matplotlib/tight_bbox.py b/lib/matplotlib/tight_bbox.py
@@ -2,6 +2,9 @@
 Helper module for the *bbox_inches* parameter in `.Figure.savefig`.
 """
 
+import contextlib
+
+from matplotlib.cbook import _setattr_cm
 from matplotlib.transforms import Bbox, TransformedBbox, Affine2D
 
 
@@ -18,46 +21,23 @@ def adjust_bbox(fig, bbox_inches, fixed_dpi=None):
     def no_op_apply_aspect(position=None):
         return
 
-    origBbox = fig.bbox
-    origBboxInches = fig.bbox_inches
-    orig_tight_layout = fig.get_tight_layout()
-    _boxout = fig.transFigure._boxout
+    stack = contextlib.ExitStack()
 
+    stack.callback(fig.set_tight_layout, fig.get_tight_layout())
     fig.set_tight_layout(False)
-    old_aspect = []
-    locator_list = []
-    sentinel = object()
+
     for ax in fig.axes:
         pos = ax.get_position(original=False).frozen()
-        locator_list.append(ax.get_axes_locator())
 
         def _l(a, r, pos=pos):
             return pos
+
+        stack.callback(ax.set_axes_locator, ax.get_axes_locator())
         ax.set_axes_locator(_l)
+
         # override the method that enforces the aspect ratio
         # on the Axes
-        if 'apply_aspect' in ax.__dict__:
-            old_aspect.append(ax.apply_aspect)
-        else:
-            old_aspect.append(sentinel)
-        ax.apply_aspect = no_op_apply_aspect
-
-    def restore_bbox():
-        for ax, loc, aspect in zip(fig.axes, locator_list, old_aspect):
-            ax.set_axes_locator(loc)
-            if aspect is sentinel:
-                # delete our no-op function which un-hides the
-                # original method
-                del ax.apply_aspect
-            else:
-                ax.apply_aspect = aspect
-
-        fig.bbox = origBbox
-        fig.bbox_inches = origBboxInches
-        fig.set_tight_layout(orig_tight_layout)
-        fig.transFigure._boxout = _boxout
-        fig.transFigure.invalidate()
-        fig.patch.set_bounds(0, 0, 1, 1)
+        stack.enter_context(_setattr_cm(ax, apply_aspect=no_op_apply_aspect))
 
     if fixed_dpi is not None:
         tr = Affine2D().scale(fixed_dpi)
@@ -68,19 +48,25 @@ def restore_bbox():
 
     _bbox = TransformedBbox(bbox_inches, tr)
 
-    fig.bbox_inches = Bbox.from_bounds(0, 0,
-                                       bbox_inches.width, bbox_inches.height)
+    stack.enter_context(
+        _setattr_cm(fig, bbox_inches=Bbox.from_bounds(
+            0, 0, bbox_inches.width, bbox_inches.height)))
     x0, y0 = _bbox.x0, _bbox.y0
     w1, h1 = fig.bbox.width * dpi_scale, fig.bbox.height * dpi_scale
-    fig.transFigure._boxout = Bbox.from_bounds(-x0, -y0, w1, h1)
+    stack.enter_context(
+        _setattr_cm(fig.transFigure,
+                    _boxout=Bbox.from_bounds(-x0, -y0, w1, h1)))
     fig.transFigure.invalidate()
+    stack.callback(fig.transFigure.invalidate)
 
-    fig.bbox = TransformedBbox(fig.bbox_inches, tr)
+    stack.enter_context(
+        _setattr_cm(fig, bbox=TransformedBbox(fig.bbox_inches, tr)))
 
+    stack.callback(fig.patch.set_bounds, 0, 0, 1, 1)
     fig.patch.set_bounds(x0 / w1, y0 / h1,
                          fig.bbox.width / w1, fig.bbox.height / h1)
 
-    return restore_bbox
+    return stack.close
 
 
 def process_figure_for_rasterizing(fig, bbox_inches_restore, fixed_dpi=None):
