Fix xkcd() not resetting context anymore.

anntzer · anntzer · commit 73bf2bbde495 · 2017-10-28T08:46:15.000-07:00
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -21,9 +21,9 @@
 import six
 
 import sys
-import warnings
 import time
 import types
+import warnings
 
 from cycler import cycler
 import matplotlib
@@ -389,28 +389,38 @@ 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
+    xkcd_ctx = 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,
+    })
+    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
@@ -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
