1# Copyright (c) 2014,2015,2017,2019 MetPy Developers.
2# Distributed under the terms of the BSD 3-Clause License.
3# SPDX-License-Identifier: BSD-3-Clause
4"""Work with custom color tables.
5
6Contains a tools for reading color tables from files, and creating instances based on a
7specific set of constraints (e.g. step size) for mapping.
8
9.. plot::
10
11   import numpy as np
12   import matplotlib.pyplot as plt
13   import metpy.plots.ctables as ctables
14
15   def plot_color_gradients(cmap_category, cmap_list, nrows):
16       fig, axes = plt.subplots(figsize=(7, 6), nrows=nrows)
17       fig.subplots_adjust(top=.93, bottom=0.01, left=0.32, right=0.99)
18       axes[0].set_title(cmap_category + ' colormaps', fontsize=14)
19
20       for ax, name in zip(axes, cmap_list):
21               ax.imshow(gradient, aspect='auto', cmap=ctables.registry.get_colortable(name))
22               pos = list(ax.get_position().bounds)
23               x_text = pos[0] - 0.01
24               y_text = pos[1] + pos[3]/2.
25               fig.text(x_text, y_text, name, va='center', ha='right', fontsize=10)
26
27       # Turn off *all* ticks & spines, not just the ones with colormaps.
28       for ax in axes:
29           ax.set_axis_off()
30
31   cmaps = list(ctables.registry)
32   cmaps = [name for name in cmaps if name[-2:]!='_r']
33   nrows = len(cmaps)
34   gradient = np.linspace(0, 1, 256)
35   gradient = np.vstack((gradient, gradient))
36
37   plot_color_gradients('MetPy', cmaps, nrows)
38   plt.show()
39"""
40
41import ast
42import logging
43from pathlib import Path
44
45import matplotlib.colors as mcolors
46
47from ..package_tools import Exporter
48
49exporter = Exporter(globals())
50
51TABLE_EXT = '.tbl'
52
53log = logging.getLogger(__name__)
54
55
56def _parse(s):
57    if hasattr(s, 'decode'):
58        s = s.decode('ascii')
59
60    if not s.startswith('#'):
61        return ast.literal_eval(s)
62
63    return None
64
65
66@exporter.export
67def read_colortable(fobj):
68    r"""Read colortable information from a file.
69
70    Reads a colortable, which consists of one color per line of the file, where
71    a color can be one of: a tuple of 3 floats, a string with a HTML color name,
72    or a string with a HTML hex color.
73
74    Parameters
75    ----------
76    fobj : a file-like object
77        A file-like object to read the colors from
78
79    Returns
80    -------
81    List of tuples
82        A list of the RGB color values, where each RGB color is a tuple of 3 floats in the
83        range of [0, 1].
84
85    """
86    ret = []
87    try:
88        for line in fobj:
89            literal = _parse(line)
90            if literal:
91                ret.append(mcolors.colorConverter.to_rgb(literal))
92        return ret
93    except (SyntaxError, ValueError) as e:
94        raise RuntimeError(f'Malformed colortable (bad line: {line})') from e
95
96
97def convert_gempak_table(infile, outfile):
98    r"""Convert a GEMPAK color table to one MetPy can read.
99
100    Reads lines from a GEMPAK-style color table file, and writes them to another file in
101    a format that MetPy can parse.
102
103    Parameters
104    ----------
105    infile : file-like object
106        The file-like object to read from
107    outfile : file-like object
108        The file-like object to write to
109
110    """
111    for line in infile:
112        if not line.startswith('!') and line.strip():
113            r, g, b = map(int, line.split())
114            outfile.write(f'({r / 255:f}, {g / 255:f}, {b / 255:f})\n')
115
116
117class ColortableRegistry(dict):
118    r"""Manages the collection of color tables.
119
120    Provides access to color tables, read collections of files, and generates
121    matplotlib's Normalize instances to go with the colortable.
122    """
123
124    def scan_resource(self, pkg, path):
125        r"""Scan a resource directory for colortable files and add them to the registry.
126
127        Parameters
128        ----------
129        pkg : str
130            The package containing the resource directory
131        path : str
132            The path to the directory with the color tables
133
134        """
135        try:
136            from importlib.resources import files as importlib_resources_files
137        except ImportError:  # Can remove when we require Python > 3.8
138            from importlib_resources import files as importlib_resources_files
139
140        for entry in (importlib_resources_files(pkg) / path).iterdir():
141            if entry.suffix == TABLE_EXT:
142                with entry.open() as stream:
143                    self.add_colortable(stream, entry.with_suffix('').name)
144
145    def scan_dir(self, path):
146        r"""Scan a directory on disk for color table files and add them to the registry.
147
148        Parameters
149        ----------
150        path : str
151            The path to the directory with the color tables
152
153        """
154        for entry in Path(path).glob('*' + TABLE_EXT):
155            if entry.is_file():
156                with entry.open() as fobj:
157                    try:
158                        self.add_colortable(fobj, entry.with_suffix('').name)
159                        log.debug('Added colortable from file: %s', entry)
160                    except RuntimeError:
161                        # If we get a file we can't handle, assume we weren't meant to.
162                        log.info('Skipping unparsable file: %s', entry)
163
164    def add_colortable(self, fobj, name):
165        r"""Add a color table from a file to the registry.
166
167        Parameters
168        ----------
169        fobj : file-like object
170            The file to read the color table from
171        name : str
172            The name under which the color table will be stored
173
174        """
175        self[name] = read_colortable(fobj)
176        self[name + '_r'] = self[name][::-1]
177
178    def get_with_steps(self, name, start, step):
179        r"""Get a color table from the registry with a corresponding norm.
180
181        Builds a `matplotlib.colors.BoundaryNorm` using `start`, `step`, and
182        the number of colors, based on the color table obtained from `name`.
183
184        Parameters
185        ----------
186        name : str
187            The name under which the color table will be stored
188        start : float
189            The starting boundary
190        step : float
191            The step between boundaries
192
193        Returns
194        -------
195        `matplotlib.colors.BoundaryNorm`, `matplotlib.colors.ListedColormap`
196            The boundary norm based on `start` and `step` with the number of colors
197            from the number of entries matching the color table, and the color table itself.
198
199        """
200        from numpy import arange
201
202        # Need one more boundary than color
203        num_steps = len(self[name]) + 1
204        boundaries = arange(start, start + step * num_steps, step)
205        return self.get_with_boundaries(name, boundaries)
206
207    def get_with_range(self, name, start, end):
208        r"""Get a color table from the registry with a corresponding norm.
209
210        Builds a `matplotlib.colors.BoundaryNorm` using `start`, `end`, and
211        the number of colors, based on the color table obtained from `name`.
212
213        Parameters
214        ----------
215        name : str
216            The name under which the color table will be stored
217        start : float
218            The starting boundary
219        end : float
220            The ending boundary
221
222        Returns
223        -------
224        `matplotlib.colors.BoundaryNorm`, `matplotlib.colors.ListedColormap`
225            The boundary norm based on `start` and `end` with the number of colors
226            from the number of entries matching the color table, and the color table itself.
227
228        """
229        from numpy import linspace
230
231        # Need one more boundary than color
232        num_steps = len(self[name]) + 1
233        boundaries = linspace(start, end, num_steps)
234        return self.get_with_boundaries(name, boundaries)
235
236    def get_with_boundaries(self, name, boundaries):
237        r"""Get a color table from the registry with a corresponding norm.
238
239        Builds a `matplotlib.colors.BoundaryNorm` using `boundaries`.
240
241        Parameters
242        ----------
243        name : str
244            The name under which the color table will be stored
245        boundaries : array_like
246            The list of boundaries for the norm
247
248        Returns
249        -------
250        `matplotlib.colors.BoundaryNorm`, `matplotlib.colors.ListedColormap`
251            The boundary norm based on `boundaries`, and the color table itself.
252
253        """
254        cmap = self.get_colortable(name)
255        return mcolors.BoundaryNorm(boundaries, cmap.N), cmap
256
257    def get_colortable(self, name):
258        r"""Get a color table from the registry.
259
260        Parameters
261        ----------
262        name : str
263            The name under which the color table will be stored
264
265        Returns
266        -------
267        `matplotlib.colors.ListedColormap`
268            The color table corresponding to `name`
269
270        """
271        return mcolors.ListedColormap(self[name], name=name)
272
273
274registry = ColortableRegistry()
275registry.scan_resource('metpy.plots', 'colortable_files')
276registry.scan_dir(Path.cwd())
277
278with exporter:
279    colortables = registry
280