Add a dedicated ColormapRegistry class

timhoffm · timhoffm · commit ce7f656d0d09 · 2021-08-02T00:48:32.000+02:00
diff --git a/.flake8 b/.flake8
@@ -45,7 +45,7 @@ per-file-ignores =
     setupext.py: E501
     tests.py: F401
 
-    lib/matplotlib/__init__.py: F401
+    lib/matplotlib/__init__.py: E401, F401
     lib/matplotlib/_api/__init__.py: F401
     lib/matplotlib/_cm.py: E122, E202, E203, E302
     lib/matplotlib/_mathtext.py: E221, E251
diff --git a/doc/api/matplotlib_configuration_api.rst b/doc/api/matplotlib_configuration_api.rst
@@ -52,6 +52,27 @@ Logging
 
 .. autofunction:: set_loglevel
 
+Colormaps
+=========
+
+.. py:data:: colormaps
+
+   A `.ColormapRegistry` instance holding all known colormaps.
+
+   Get a `.Colormap` using item access::
+
+        import matplotlib as mpl
+        cmap = mpl.colormaps['viridis']
+
+   Returned `.Colormap`\s are copies, so that their modification does not
+   change the global definition of the colormap.
+
+   Additional colormaps can be added via `.ColormapRegistry.register`::
+
+        mpl.colormaps.register(my_colormap)
+
+.. autodata:: colormaps
+
 Miscellaneous
 =============
 
diff --git a/doc/users/next_whats_new/colormaps.rst b/doc/users/next_whats_new/colormaps.rst
@@ -0,0 +1,18 @@
+Colormaps registry
+------------------
+
+Colormaps are now managed via `matplotlib.colormaps`, which is a
+`.ColormapRegistry`.
+
+Colormaps can be obtained using item access::
+
+    import matplotlib as mpl
+    cmap = mpl.colormaps['viridis']
+
+To register new colormaps use::
+
+    mpl.colormaps.register(my_colormap)
+
+The use of `matplotlib.cm.get_cmap` and `matplotlib.cm.register_cmap` is
+discouraged in favor of the above. Within `.pyplot` the use of
+``plt.get_cmap()`` and ``plt.register_cmap()`` will continue to be supported.
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1445,3 +1445,8 @@ def inner(ax, *args, data=None, **kwargs):
 _log.debug('interactive is %s', is_interactive())
 _log.debug('platform is %s', sys.platform)
 _log.debug('loaded modules: %s', list(sys.modules))
+
+
+# workaround: we must defer colormaps import to after loading rcParams, becuase
+# colormap creation depends on rcParams
+from matplotlib.cm import _colormaps as colormaps
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -15,7 +15,7 @@
   normalization.
 """
 
-from collections.abc import MutableMapping
+from collections.abc import Mapping, MutableMapping
 
 import numpy as np
 from numpy import ma
@@ -91,13 +91,86 @@ def _warn_deprecated(self):
         )
 
 
+class ColormapRegistry(Mapping):
+    r"""
+    Container for colormaps that are known to Matplotlib by name.
+
+    There should be no need for users to instantiate this class themselves.
+
+    The universal registry instance is `matplotlib.colormaps`.
+
+    Read access uses a dict-like interface mapping names to `.Colormap`\s::
+
+        import matplotlib as mpl
+        cmap = mpl.colormaps['viridis']
+
+    Returned `.Colormap`\s are copies, so that their modification does not
+    change the global definition of the colormap.
+
+    Additional colormaps can be added via `.ColormapRegistry.register`::
+
+        mpl.colormaps.register(my_colormap)
+    """
+    def __init__(self, cmaps):
+        self._cmaps = cmaps
+
+    def __getitem__(self, item):
+        try:
+            return self._cmaps[item].copy()
+        except KeyError:
+            raise KeyError(f"{item!r} is not a known colormap name")
+
+    def __iter__(self):
+        return iter(self._cmaps)
+
+    def __len__(self):
+        return len(self._cmaps)
+
+    def __str__(self):
+        return ('ColormapRegistry; available colormaps:\n' +
+                ', '.join(f"'{name}'" for name in self))
+
+    def register(self, cmap, *, name=None, force=False):
+        """
+        Register a new colormap.
+
+        The colormap name can then be used as a string argument to any ``cmap``
+        parameter in Matplotlib. It is also available in `.pyplot.get_cmap`.
+
+        The colormap registry stores a copy of the given colormap, so that
+        future changes to the original colormap instance do not affect the
+        registered colormap. Think of this as the registry taking a snapshot
+        of the colormap at registration.
+
+        Parameters
+        ----------
+        cmap : matplotlib.colors.Colormap
+            The colormap to register.
+
+        name : str, optional
+            The name for the colormap. If not given, the
+            :attr:`~.Colormap.name` attribute of *cmap* is used.
+
+        force: bool, default: False
+
+        """
+        name_ = name or cmap.name
+        if name_ in self and not force:
+            raise ValueError(
+                f'A colormap named "{name}" is already registered.')
+        register_cmap(name, cmap.copy())
+
+
 _cmap_registry = _gen_cmap_registry()
 globals().update(_cmap_registry)
 # This is no longer considered public API
 cmap_d = _DeprecatedCmapDictWrapper(_cmap_registry)
 __builtin_cmaps = tuple(_cmap_registry)
 
-# Continue with definitions ...
+# public acces to the colormaps should be via `matplotlib.colormaps`. For now,
+# we still create the registry here, but that should stay an implementation
+# detail.
+_colormaps = ColormapRegistry(_cmap_registry)
 
 
 def register_cmap(name=None, cmap=None, *, override_builtin=False):
