Merge pull request #19369 from anntzer/cmset

QuLogic · web-flow · commit cb756d3a57a5 · 2021-05-05T20:07:28.000-04:00
Add Artist._cm_set for temporarily setting an Artist property.
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -1,4 +1,5 @@
 from collections import namedtuple
+import contextlib
 from functools import wraps
 import inspect
 import logging
@@ -1154,6 +1155,18 @@ def set(self, **kwargs):
         kwargs = cbook.normalize_kwargs(kwargs, self)
         return self.update(kwargs)
 
+    @contextlib.contextmanager
+    def _cm_set(self, **kwargs):
+        """
+        `.Artist.set` context-manager that restores original values at exit.
+        """
+        orig_vals = {k: getattr(self, f"get_{k}")() for k in kwargs}
+        try:
+            self.set(**kwargs)
+            yield
+        finally:
+            self.set(**orig_vals)
+
     def findobj(self, match=None, include_self=True):
         """
         Find artist objects.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -3122,8 +3122,7 @@ def redraw_in_frame(self):
         with ExitStack() as stack:
             for artist in [*self._get_axis_list(),
                            self.title, self._left_title, self._right_title]:
-                stack.callback(artist.set_visible, artist.get_visible())
-                artist.set_visible(False)
+                stack.enter_context(artist._cm_set(visible=False))
             self.draw(self.figure._cachedRenderer)
 
     def get_renderer_cache(self):
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -26,7 +26,7 @@
 """
 
 from collections import namedtuple
-from contextlib import contextmanager, nullcontext
+from contextlib import ExitStack, contextmanager, nullcontext
 from enum import Enum, IntEnum
 import functools
 import importlib
@@ -2256,22 +2256,16 @@ def print_figure(
 
         # Remove the figure manager, if any, to avoid resizing the GUI widget.
         with cbook._setattr_cm(self, manager=None), \
-                cbook._setattr_cm(self.figure, dpi=dpi), \
-                cbook._setattr_cm(canvas, _is_saving=True):
-            origfacecolor = self.figure.get_facecolor()
-            origedgecolor = self.figure.get_edgecolor()
-
-            if facecolor is None:
-                facecolor = rcParams['savefig.facecolor']
-            if cbook._str_equal(facecolor, 'auto'):
-                facecolor = origfacecolor
-            if edgecolor is None:
-                edgecolor = rcParams['savefig.edgecolor']
-            if cbook._str_equal(edgecolor, 'auto'):
-                edgecolor = origedgecolor
-
-            self.figure.set_facecolor(facecolor)
-            self.figure.set_edgecolor(edgecolor)
+             cbook._setattr_cm(self.figure, dpi=dpi), \
+             cbook._setattr_cm(canvas, _is_saving=True), \
+             ExitStack() as stack:
+
+            for prop in ["facecolor", "edgecolor"]:
+                color = locals()[prop]
+                if color is None:
+                    color = rcParams[f"savefig.{prop}"]
+                if not cbook._str_equal(color, "auto"):
+                    stack.enter_context(self.figure._cm_set(**{prop: color}))
 
             if bbox_inches is None:
                 bbox_inches = rcParams['savefig.bbox']
@@ -2306,8 +2300,7 @@ def print_figure(
                 _bbox_inches_restore = None
 
             # we have already done CL above, so turn it off:
-            cl_state = self.figure.get_constrained_layout()
-            self.figure.set_constrained_layout(False)
+            stack.enter_context(self.figure._cm_set(constrained_layout=False))
             try:
                 # _get_renderer may change the figure dpi (as vector formats
                 # force the figure dpi to 72), so we need to set it again here.
@@ -2323,11 +2316,7 @@ def print_figure(
                 if bbox_inches and restore_bbox:
                     restore_bbox()
 
-                self.figure.set_facecolor(origfacecolor)
-                self.figure.set_edgecolor(origedgecolor)
                 self.figure.set_canvas(self)
-            # reset to cached state
-            self.figure.set_constrained_layout(cl_state)
             return result
 
     @classmethod
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -14,6 +14,7 @@
     Control the default spacing between subplots.
 """
 
+from contextlib import ExitStack
 import inspect
 import logging
 from numbers import Integral
@@ -2936,23 +2937,15 @@ def savefig(self, fname, *, transparent=None, **kwargs):
         if transparent is None:
             transparent = mpl.rcParams['savefig.transparent']
 
-        if transparent:
-            kwargs.setdefault('facecolor', 'none')
-            kwargs.setdefault('edgecolor', 'none')
-            original_axes_colors = []
-            for ax in self.axes:
-                patch = ax.patch
-                original_axes_colors.append((patch.get_facecolor(),
-                                             patch.get_edgecolor()))
-                patch.set_facecolor('none')
-                patch.set_edgecolor('none')
-
-        self.canvas.print_figure(fname, **kwargs)
-
-        if transparent:
-            for ax, cc in zip(self.axes, original_axes_colors):
-                ax.patch.set_facecolor(cc[0])
-                ax.patch.set_edgecolor(cc[1])
+        with ExitStack() as stack:
+            if transparent:
+                kwargs.setdefault('facecolor', 'none')
+                kwargs.setdefault('edgecolor', 'none')
+                for ax in self.axes:
+                    stack.enter_context(
+                        ax.patch._cm_set(facecolor='none', edgecolor='none'))
+
+            self.canvas.print_figure(fname, **kwargs)
 
     def ginput(self, n=1, timeout=30, show_clicks=True,
                mouse_add=MouseButton.LEFT,
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -1735,8 +1735,7 @@ def update_background(self, event):
         with ExitStack() as stack:
             if needs_redraw:
                 for artist in self.artists:
-                    stack.callback(artist.set_visible, artist.get_visible())
-                    artist.set_visible(False)
+                    stack.enter_context(artist._cm_set(visible=False))
                 self.canvas.draw()
             self.background = self.canvas.copy_from_bbox(self.ax.bbox)
         if needs_redraw:
