Add style context manager and tests

tonysyu · tonysyu · commit 7769b2940768 · 2013-11-17T14:04:45.000-06:00
diff --git a/lib/matplotlib/style/core.py b/lib/matplotlib/style/core.py
@@ -3,26 +3,30 @@
 
 ``use``
     Select style sheet to override the current matplotlib settings.
+``context``
+    Context manager to use a style sheet temporarily.
 ``available``
     List available style sheets.
 ``library``
     A dictionary of style names and matplotlib settings.
 """
 import os
 import re
+import contextlib
 
 import numpy as np
 import matplotlib as mpl
 
 
-__all__ = ['use', 'available', 'library']
+__all__ = ['use', 'context', 'available', 'library', 'reload_library']
 
 
 _here = os.path.abspath(os.path.dirname(__file__))
 BASE_LIBRARY_PATH = os.path.join(_here, 'stylelib')
 # Users may want multiple library paths, so store a list of paths.
 USER_LIBRARY_PATHS = [os.path.join('~', '.matplotlib', 'stylelib')]
-STYLE_FILE_PATTERN = re.compile('([\S]+).style$')
+STYLE_EXTENSION = 'style'
+STYLE_FILE_PATTERN = re.compile('([\S]+).%s$' % STYLE_EXTENSION)
 
 
 def is_style_file(filename):
@@ -31,7 +35,7 @@ def is_style_file(filename):
 
 
 def use(name):
-    """Use matplotlib rc parameters from a pre-defined name or from a file.
+    """Use matplotlib style settings from a known style sheet or from a file.
 
     Parameters
     ----------
@@ -55,6 +59,28 @@ def use(name):
             mpl.rcParams.update(library[style])
 
 
+@contextlib.contextmanager
+def context(name, after_reset=False):
+    """Context manager for using style settings temporarily.
+
+    Parameters
+    ----------
+    name : str or list of str
+        Name of style or path/URL to a style file. For a list of available
+        style names, see `style.available`. If given a list, each style is
+        applied from first to last in the list.
+    after_reset : bool
+        If True, apply style after resetting settings to their defaults;
+        otherwise, apply style on top of the current settings.
+    """
+    initial_settings = mpl.rcParams.copy()
+    if after_reset:
+        mpl.rcdefaults()
+    use(name)
+    yield
+    mpl.rcParams.update(initial_settings)
+
+
 def load_base_library():
     """Load style library defined in this package."""
     library = dict()
@@ -114,5 +140,13 @@ def update_nested_dict(main_dict, new_dict):
 # Load style library
 # ==================
 _base_library = load_base_library()
-library = update_user_library(_base_library)
-available = library.keys()
+
+library = None
+available = []
+
+def reload_library():
+    """Reload style library."""
+    global library, available
+    library = update_user_library(_base_library)
+    available[:] = library.keys()
+reload_library()
diff --git a/lib/matplotlib/style/tests/test_core.py b/lib/matplotlib/style/tests/test_core.py
@@ -0,0 +1,61 @@
+import os
+import shutil
+import tempfile
+from contextlib import contextmanager
+
+import matplotlib as mpl
+from matplotlib import style
+from matplotlib.style.core import USER_LIBRARY_PATHS, STYLE_EXTENSION
+
+
+PARAM = 'image.cmap'
+VALUE = 'pink'
+DUMMY_SETTINGS = {PARAM: VALUE}
+
+
+@contextmanager
+def temp_style(style_name, settings=None):
+    """Context manager to create a style sheet in a temporary directory."""
+    settings = DUMMY_SETTINGS
+    temp_file = '%s.%s' % (style_name, STYLE_EXTENSION)
+
+    # Write style settings to file in the temp directory.
+    tempdir = tempfile.mkdtemp()
+    with open(os.path.join(tempdir, temp_file), 'w') as f:
+        for k, v in settings.iteritems():
+            f.write('%s: %s' % (k, v))
+
+    # Add temp directory to style path and reload so we can access this style.
+    USER_LIBRARY_PATHS.append(tempdir)
+    style.reload_library()
+
+    try:
+        yield
+    finally:
+        shutil.rmtree(tempdir)
+
+
+def test_available():
+    with temp_style('_test_', DUMMY_SETTINGS):
+        assert '_test_' in style.available
+
+
+def test_use():
+    mpl.rcParams[PARAM] = 'gray'
+    with temp_style('test', DUMMY_SETTINGS):
+        style.use('test')
+        assert mpl.rcParams[PARAM] == VALUE
+
+
+def test_context():
+    mpl.rcParams[PARAM] = 'gray'
+    with temp_style('test', DUMMY_SETTINGS):
+        with style.context('test'):
+            assert mpl.rcParams[PARAM] == VALUE
+    # Check that this value is reset after the exiting the context.
+    assert mpl.rcParams[PARAM] == 'gray'
+
+
+if __name__ == '__main__':
+    from numpy import testing
+    testing.run_module_suite()
