Merge pull request #15303 from anntzer/rc_ctxd

timhoffm · web-flow · commit 63b63128d189 · 2020-03-24T19:02:07.000+01:00
Make it possible to use rc_context as a decorator.
diff --git a/doc/users/next_whats_new/2018-09-19-AL.rst b/doc/users/next_whats_new/2018-09-19-AL.rst
@@ -0,0 +1,11 @@
+:orphan:
+
+`matplotlib.rc_context` is now a `contextlib.contextmanager`
+````````````````````````````````````````````````````````````
+
+`matplotlib.rc_context` can now be used as a decorator (technically, it is now
+implemented as a `contextlib.contextmanager`), e.g. ::
+
+    @rc_context({"lines.linewidth": 2})
+    def some_function(...):
+        ...
diff --git a/doc/users/prev_whats_new/whats_new_1.3.rst b/doc/users/prev_whats_new/whats_new_1.3.rst
@@ -86,7 +86,7 @@ New plotting features
 To give your plots a sense of authority that they may be missing,
 Michael Droettboom (inspired by the work of many others in
 :ghpull:`1329`) has added an `xkcd-style <http://xkcd.com/>`__ sketch
-plotting mode.  To use it, simply call :func:`matplotlib.pyplot.xkcd`
+plotting mode.  To use it, simply call `matplotlib.pyplot.xkcd`
 before creating your plot. For really fine control, it is also possible
 to modify each artist's sketch parameters individually with
 :meth:`matplotlib.artist.Artist.set_sketch_params`.
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1035,7 +1035,8 @@ def rc_file(fname, *, use_default_template=True):
                          if k not in STYLE_BLACKLIST})
 
 
-class rc_context:
+@contextlib.contextmanager
+def rc_context(rc=None, fname=None):
     """
     Return a context manager for managing rc settings.
 
@@ -1052,49 +1053,24 @@ class rc_context:
         with mpl.rc_context(rc={'text.usetex': True}, fname='screen.rc'):
             plt.plot(x, a)
 
-    The 'rc' dictionary takes precedence over the settings loaded from
-    'fname'.  Passing a dictionary only is also valid. For example a
-    common usage is::
+    The *rc* dictionary takes precedence over the settings loaded from *fname*.
+    Passing a dictionary only is also valid.  For example, a common usage is::
 
-        with mpl.rc_context(rc={'interactive': False}):
+        with mpl.rc_context({'interactive': False}):
             fig, ax = plt.subplots()
             ax.plot(range(3), range(3))
             fig.savefig('A.png', format='png')
             plt.close(fig)
     """
-    # While it may seem natural to implement rc_context using
-    # contextlib.contextmanager, that would entail always calling the finally:
-    # clause of the contextmanager (which restores the original rcs) including
-    # during garbage collection; as a result, something like `plt.xkcd();
-    # gc.collect()` would result in the style being lost (as `xkcd()` is
-    # implemented on top of rc_context, and nothing is holding onto context
-    # manager except possibly circular references.
-
-    def __init__(self, rc=None, fname=None):
-        self._orig = rcParams.copy()
-        try:
-            if fname:
-                rc_file(fname)
-            if rc:
-                rcParams.update(rc)
-        except Exception:
-            self.__fallback()
-            raise
-
-    def __fallback(self):
-        # If anything goes wrong, revert to the original rcs.
-        updated_backend = self._orig['backend']
-        dict.update(rcParams, self._orig)
-        # except for the backend.  If the context block triggered resolving
-        # the auto backend resolution keep that value around
-        if self._orig['backend'] is rcsetup._auto_backend_sentinel:
-            rcParams['backend'] = updated_backend
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, exc_tb):
-        self.__fallback()
+    orig = rcParams.copy()
+    try:
+        if fname:
+            rc_file(fname)
+        if rc:
+            rcParams.update(rc)
+        yield
+    finally:
+        dict.update(rcParams, orig)  # Revert to the original rcs.
 
 
 def use(backend, *, force=True):
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -408,12 +408,11 @@ def setp(obj, *args, **kwargs):
 
 def xkcd(scale=1, length=100, randomness=2):
     """
-    Turn on `xkcd <https://xkcd.com/>`_ sketch-style drawing mode.
-    This will only have effect on things drawn after this function is
-    called.
+    Turn on `xkcd <https://xkcd.com/>`_ sketch-style drawing mode.  This will
+    only have effect on things drawn after this function is called.
 
     For best results, the "Humor Sans" font should be installed: it is
-    not included with matplotlib.
+    not included with Matplotlib.
 
     Parameters
     ----------
@@ -440,29 +439,46 @@ def xkcd(scale=1, length=100, randomness=2):
         # This figure will be in regular style
         fig2 = plt.figure()
     """
-    if rcParams['text.usetex']:
-        raise RuntimeError(
-            "xkcd mode is not compatible with text.usetex = True")
-
-    from matplotlib import patheffects
-    return rc_context({
-        'font.family': ['xkcd', 'xkcd Script', 'Humor Sans', 'Comic Neue',
-                        'Comic Sans MS'],
-        'font.size': 14.0,
-        'path.sketch': (scale, length, randomness),
-        'path.effects': [patheffects.withStroke(linewidth=4, foreground="w")],
-        'axes.linewidth': 1.5,
-        'lines.linewidth': 2.0,
-        'figure.facecolor': 'white',
-        'grid.linewidth': 0.0,
-        'axes.grid': False,
-        'axes.unicode_minus': False,
-        'axes.edgecolor': 'black',
-        'xtick.major.size': 8,
-        'xtick.major.width': 3,
-        'ytick.major.size': 8,
-        'ytick.major.width': 3,
-    })
+    return _xkcd(scale, length, randomness)
+
+
+class _xkcd:
+    # This cannot be implemented in terms of rc_context() because this needs to
+    # work as a non-contextmanager too.
+
+    def __init__(self, scale, length, randomness):
+        self._orig = rcParams.copy()
+
+        if rcParams['text.usetex']:
+            raise RuntimeError(
+                "xkcd mode is not compatible with text.usetex = True")
+
+        from matplotlib import patheffects
+        rcParams.update({
+            'font.family': ['xkcd', 'xkcd Script', 'Humor Sans', 'Comic Neue',
+                            'Comic Sans MS'],
+            'font.size': 14.0,
+            'path.sketch': (scale, length, randomness),
+            'path.effects': [
+                patheffects.withStroke(linewidth=4, foreground="w")],
+            'axes.linewidth': 1.5,
+            'lines.linewidth': 2.0,
+            'figure.facecolor': 'white',
+            'grid.linewidth': 0.0,
+            'axes.grid': False,
+            'axes.unicode_minus': False,
+            'axes.edgecolor': 'black',
+            'xtick.major.size': 8,
+            'xtick.major.width': 3,
+            'ytick.major.size': 8,
+            'ytick.major.width': 3,
+        })
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        dict.update(rcParams, self._orig)
 
 
 ## Figures ##
diff --git a/lib/matplotlib/tests/test_rcparams.py b/lib/matplotlib/tests/test_rcparams.py
@@ -54,6 +54,14 @@ def test_rcparams(tmpdir):
         assert mpl.rcParams['lines.linewidth'] == 44
     assert mpl.rcParams['lines.linewidth'] == linewidth
 
+    # test context as decorator (and test reusability, by calling func twice)
+    @mpl.rc_context({'lines.linewidth': 44})
+    def func():
+        assert mpl.rcParams['lines.linewidth'] == 44
+
+    func()
+    func()
+
     # test rc_file
     mpl.rc_file(rcpath)
     assert mpl.rcParams['lines.linewidth'] == 33
