Fix xkcd() not resetting context anymore.

anntzer · anntzer · commit a02949eb1b08 · 2017-10-28T08:24:28.000-07:00
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1222,16 +1222,17 @@ def rc_context(rc=None, fname=None):
 
     """
 
-    orig = rcParams.copy()
+    _rcParams = rcParams  # Prevent it from being set to None at shutdown.
+    orig = _rcParams.copy()
     try:
         if fname:
             rc_file(fname)
         if rc:
-            rcParams.update(rc)
+            _rcParams.update(rc)
         yield
     finally:
         # No need to revalidate the original values.
-        dict.update(rcParams, orig)
+        dict.update(_rcParams, orig)
 
 
 _use_error_msg = """
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -21,9 +21,13 @@
 import six
 
 import sys
-import warnings
 import time
 import types
+import warnings
+try:
+    from contextlib import ExitStack
+except ImportError:
+    from contextlib2 import ExitStack  # Py2
 
 from cycler import cycler
 import matplotlib
@@ -389,28 +393,30 @@ def xkcd(scale=1, length=100, randomness=2):
             "xkcd mode is not compatible with text.usetex = True")
 
     from matplotlib import patheffects
-    context = rc_context()
-    try:
-        rcParams['font.family'] = ['xkcd', 'Humor Sans', 'Comic Sans MS']
-        rcParams['font.size'] = 14.0
-        rcParams['path.sketch'] = (scale, length, randomness)
-        rcParams['path.effects'] = [
-            patheffects.withStroke(linewidth=4, foreground="w")]
-        rcParams['axes.linewidth'] = 1.5
-        rcParams['lines.linewidth'] = 2.0
-        rcParams['figure.facecolor'] = 'white'
-        rcParams['grid.linewidth'] = 0.0
-        rcParams['axes.grid'] = False
-        rcParams['axes.unicode_minus'] = False
-        rcParams['axes.edgecolor'] = 'black'
-        rcParams['xtick.major.size'] = 8
-        rcParams['xtick.major.width'] = 3
-        rcParams['ytick.major.size'] = 8
-        rcParams['ytick.major.width'] = 3
-    except:
-        context.__exit__(*sys.exc_info())
-        raise
-    return context
+    # 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
+    # (implemented with ExitStack) 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.
+    xkcd._stack = stack = ExitStack()
+    stack.enter_context(rc_context({
+        'font.family': ['xkcd', 'Humor Sans', '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 stack
 
 
 ## Figures ##
diff --git a/lib/matplotlib/tests/test_style.py b/lib/matplotlib/tests/test_style.py
@@ -11,7 +11,7 @@
 import pytest
 
 import matplotlib as mpl
-from matplotlib import style
+from matplotlib import pyplot as plt, style
 from matplotlib.style.core import USER_LIBRARY_PATHS, STYLE_EXTENSION
 
 import six
@@ -158,3 +158,16 @@ def test_alias(equiv_styles):
     rc_base = rc_dicts[0]
     for nm, rc in zip(equiv_styles[1:], rc_dicts[1:]):
         assert rc_base == rc
+
+
+def test_xkcd_no_cm():
+    assert mpl.rcParams["path.sketch"] is None
+    plt.xkcd()
+    assert mpl.rcParams["path.sketch"] == (1, 100, 2)
+
+
+def test_xkcd_cm():
+    assert mpl.rcParams["path.sketch"] is None
+    with plt.xkcd():
+        assert mpl.rcParams["path.sketch"] == (1, 100, 2)
+    assert mpl.rcParams["path.sketch"] is None
diff --git a/setupext.py b/setupext.py
@@ -1465,7 +1465,10 @@ def get_install_requires(self):
             "six>=1.10",
         ]
         if sys.version_info < (3,):
-            install_requires += ["backports.functools_lru_cache"]
+            install_requires += [
+                "backports.functools_lru_cache",
+                "contextlib2>=0.5.0",
+            ]
         if sys.version_info < (3,) and os.name == "posix":
             install_requires += ["subprocess32"]
         return install_requires
