Merge pull request #10856 from anntzer/xkcdgc

timhoffm · web-flow · commit 78224ef2da5e · 2018-03-21T23:10:38.000+01:00
Fix xkcd style garbage collection.
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1243,8 +1243,7 @@ def rc_file(fname):
     rcParams.update(rc_params_from_file(fname))
 
 
-@contextlib.contextmanager
-def rc_context(rc=None, fname=None):
+class rc_context:
     """
     Return a context manager for managing rc settings.
 
@@ -1274,19 +1273,33 @@ def rc_context(rc=None, fname=None):
             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:
+            # If anything goes wrong, revert to the original rcs.
+            dict.update(rcParams, self._orig)
+            raise
 
-    orig = rcParams.copy()
-    try:
-        if fname:
-            rc_file(fname)
-        if rc:
-            rcParams.update(rc)
-        yield
-    finally:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
         # No need to revalidate the original values.
-        dict.update(rcParams, orig)
+        dict.update(rcParams, self._orig)
 
 
 _use_error_msg = """
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -387,7 +387,7 @@ def xkcd(scale=1, length=100, randomness=2):
             "xkcd mode is not compatible with text.usetex = True")
 
     from matplotlib import patheffects
-    xkcd_ctx = rc_context({
+    return rc_context({
         'font.family': ['xkcd', 'Humor Sans', 'Comic Sans MS'],
         'font.size': 14.0,
         'path.sketch': (scale, length, randomness),
@@ -404,21 +404,6 @@ def xkcd(scale=1, length=100, randomness=2):
         'ytick.major.size': 8,
         'ytick.major.width': 3,
     })
-    xkcd_ctx.__enter__()
-
-    # In order to make the call to `xkcd` that does not use a context manager
-    # (cm) work, we need to enter into the cm ourselves, and return a dummy
-    # cm that does nothing on entry and cleans up the xkcd context on exit.
-    # Additionally, we need to keep a reference to the dummy cm because it
-    # would otherwise be exited when GC'd.
-
-    class dummy_ctx(object):
-        def __enter__(self):
-            pass
-
-        __exit__ = xkcd_ctx.__exit__
-
-    return dummy_ctx()
 
 
 ## Figures ##
diff --git a/lib/matplotlib/tests/test_style.py b/lib/matplotlib/tests/test_style.py
@@ -1,11 +1,12 @@
 from __future__ import absolute_import, division, print_function
 
+from collections import OrderedDict
+from contextlib import contextmanager
+import gc
 import os
 import shutil
 import tempfile
 import warnings
-from collections import OrderedDict
-from contextlib import contextmanager
 
 import pytest
 
@@ -163,6 +164,8 @@ def test_xkcd_no_cm():
     assert mpl.rcParams["path.sketch"] is None
     plt.xkcd()
     assert mpl.rcParams["path.sketch"] == (1, 100, 2)
+    gc.collect()
+    assert mpl.rcParams["path.sketch"] == (1, 100, 2)
 
 
 def test_xkcd_cm():
