Merge pull request #11723 from anntzer/picklablecmaps

dstansby · web-flow · commit 3448de363d66 · 2018-07-21T14:09:50.000-07:00
Start work on making colormaps picklable.
diff --git a/.flake8 b/.flake8
@@ -28,7 +28,7 @@ per-file-ignores =
     tools/compare_backend_driver_results.py: E501
     tools/subset.py: E221, E231, E251, E261, E302, E501, E701
 
-    matplotlib/_cm.py: E202, E203
+    matplotlib/_cm.py: E202, E203, E302
     matplotlib/_mathtext_data.py: E203, E261
     matplotlib/backend_bases.py: E225
     matplotlib/backends/_backend_tk.py: E203, E222, E225, E231, E271, E301, E303, E401, E501, E701
diff --git a/lib/matplotlib/_cm.py b/lib/matplotlib/_cm.py
@@ -41,17 +41,17 @@
                 'blue':  ((0., 0., 0.),
                           (1.0, 0.4975, 0.4975))}
 
-_flag_data = {
-        'red': lambda x: 0.75 * np.sin((x * 31.5 + 0.25) * np.pi) + 0.5,
-        'green': lambda x: np.sin(x * 31.5 * np.pi),
-        'blue': lambda x: 0.75 * np.sin((x * 31.5 - 0.25) * np.pi) + 0.5,
-}
 
-_prism_data = {
-        'red': lambda x: 0.75 * np.sin((x * 20.9 + 0.25) * np.pi) + 0.67,
-        'green': lambda x: 0.75 * np.sin((x * 20.9 - 0.25) * np.pi) + 0.33,
-        'blue': lambda x: -1.1 * np.sin((x * 20.9) * np.pi),
-}
+def _flag_red(x): return 0.75 * np.sin((x * 31.5 + 0.25) * np.pi) + 0.5
+def _flag_green(x): return np.sin(x * 31.5 * np.pi)
+def _flag_blue(x): return 0.75 * np.sin((x * 31.5 - 0.25) * np.pi) + 0.5
+def _prism_red(x): return 0.75 * np.sin((x * 20.9 + 0.25) * np.pi) + 0.67
+def _prism_green(x): return 0.75 * np.sin((x * 20.9 - 0.25) * np.pi) + 0.33
+def _prism_blue(x): return -1.1 * np.sin((x * 20.9) * np.pi)
+
+
+_flag_data = {'red': _flag_red, 'green': _flag_green, 'blue': _flag_blue}
+_prism_data = {'red': _prism_red, 'green': _prism_green, 'blue': _prism_blue}
 
 
 def cubehelix(gamma=1.0, s=0.5, r=-1.5, h=1.0):
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -5,6 +5,8 @@
 See :doc:`/tutorials/colors/colormaps` for an in-depth discussion of colormaps.
 """
 
+import functools
+
 import numpy as np
 from numpy import ma
 
@@ -22,10 +24,12 @@
 # reversed colormaps have '_r' appended to the name.
 
 
-def _reverser(f):
-    def freversed(x):
-        return f(1 - x)
-    return freversed
+def _reverser(f, x=None):
+    """Helper such that ``_reverser(f)(x) == f(1 - x)``."""
+    if x is None:
+        # Returning a partial object keeps it picklable.
+        return functools.partial(_reverser, f)
+    return f(1 - x)
 
 
 def revcmap(data):
