Fix xkcd style garbage collection.

anntzer · anntzer · commit 04e382fd331a · 2018-03-21T08:49:35.000-07:00
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1246,8 +1246,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.
 
@@ -1277,19 +1276,28 @@ def rc_context(rc=None, fname=None):
             ax.plot(range(3), range(3))
             fig.savefig('A.png', format='png')
             plt.close(fig)
-
     """
-
-    orig = rcParams.copy()
-    try:
+    # 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()
         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():
