from functools import wraps

import brewer2mpl
import numpy as np
import matplotlib as mpl
from matplotlib import cm



# Get Set2 from ColorBrewer, a set of colors deemed colorblind-safe and
# pleasant to look at by Drs. Cynthia Brewer and Mark Harrower of Pennsylvania
# State University. These colors look lovely together, and are less
# saturated than those colors in Set1. For more on ColorBrewer, see:
# - Flash-based interactive map:
#     http://colorbrewer2.org/
# - A quick visual reference to every ColorBrewer scale:
#     http://bl.ocks.org/mbostock/5577023

#class Common(object):
#    def __init__(self):
set2 = brewer2mpl.get_map('Set2', 'qualitative', 8).mpl_colors

# Another ColorBrewer scale. This one has nice "traditional" colors like
# reds and blues
set1 = brewer2mpl.get_map('Set1', 'qualitative', 9).mpl_colors

# A colormapcycle for stacked barplots
stackmaps = [brewer2mpl.get_map('YlGn', 'sequential', 8).mpl_colormap,
             brewer2mpl.get_map('YlOrRd', 'sequential', 8).mpl_colormap]

# This decorator makes it possible to change the color cycle inside
# prettyplotlib  without affecting pyplot
def pretty(func):
    rcParams = {'axes.color_cycle': set2, 'lines.linewidth': .75}

    @wraps(func)
    def wrapper(*args, **kwargs):
        with mpl.rc_context(rc=rcParams):
            return func(*args, **kwargs)
    return wrapper

# This function returns a colorlist for barplots
def getcolors(cmap, yvals, n):
    if isinstance(cmap, bool):
        cmap = stackmaps[n%len(stackmaps)]
    return [cmap( abs(int((float(yval)/np.max(yvals))*cmap.N) )) for yval in yvals]


# Set some commonly used colors
almost_black = '#262626'
light_grey = np.array([float(248) / float(255)] * 3)

reds = cm.Reds
reds.set_bad('white')
reds.set_under('white')

blues_r = cm.Blues_r
blues_r.set_bad('white')
blues_r.set_under('white')

# Need to 'reverse' red to blue so that blue=cold=small numbers,
# and red=hot=large numbers with '_r' suffix
blue_red = brewer2mpl.get_map('RdBu', 'Diverging', 11,
                              reverse=True).mpl_colormap

@@ -1022,16 +1022,24 @@
 %0A# This 
+context-
 decorato
@@ -1107,17 +1107,16 @@
 plotlib 
- 
 without 
@@ -1136,24 +1136,21 @@
 lot%0A
-def pretty(func)
+class _pretty
 :%0A  
@@ -1213,17 +1213,73 @@
 ': .75%7D%0A
-%0A
+    mpl_contexts = %5B%5D%0A%0A    def __call__(self, func):%0A    
     @wra
@@ -1287,16 +1287,20 @@
 s(func)%0A
+    
     def 
@@ -1333,20 +1333,142 @@
         
+    
 with
+ self:%0A                return func(*args, **kwargs)%0A        return wrapper%0A%0A    def __enter__(self):%0A        context =
  mpl.rc_
@@ -1478,16 +1478,21 @@
 text(rc=
+self.
 rcParams
@@ -1492,17 +1492,16 @@
 cParams)
-:
 %0A       
@@ -1505,59 +1505,183 @@
     
-    return func(*args, **kw
+self.mpl_contexts.append(context)%0A        return context.__enter__()%0A%0A    def __exit__(self, *
 args)
-%0A
+:%0A 
     
+   
 return 
-wrapper
+self.mpl_contexts.pop().__exit__(*args)%0A%0Apretty = _pretty()
 %0A%0A# 

Make ppl.pretty a context-decorator.