cmaps.template

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import re
from glob import glob
from packaging import version

import matplotlib
import matplotlib.cm
import numpy as np

from ._version import __version__
from .colormap import Colormap

CMAPSFILE_DIR = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), 'colormaps')
USER_CMAPFILE_DIR = os.environ.get('CMAP_DIR')


if version.parse(matplotlib.__version__) < version.parse('3.2.0'):
    raise Exception('cmaps of version {} only supports matplotlib greater than 3.2'.format(__version__))

if version.parse(matplotlib.__version__) >= version.parse('3.7'):
    get_cmap = matplotlib.colormaps.get_cmap
    register_cmap = matplotlib.colormaps.register
else:
    get_cmap = matplotlib.cm.get_cmap
    register_cmap = matplotlib.cm.register_cmap
    

class Cmaps(object):
    """colormaps"""

    def __init__(self, ):
        self._parse_cmaps()
        self.__version__ = __version__

    def _coltbl(self, cmap_file):
        pattern = re.compile(r'(\d\.?\d*)\s+(\d\.?\d*)\s+(\d\.?\d*).*')
        with open(cmap_file) as cmap:
            cmap_buff = cmap.read()
        cmap_buff = re.compile('ncolors.*\n').sub('', cmap_buff)
        if re.search(r'\s*\d\.\d*', cmap_buff):
            return np.asarray(pattern.findall(cmap_buff), 'f4')
        else:
            return np.asarray(pattern.findall(cmap_buff), 'u1') / 255.


    def _parse_cmaps(self):
        if USER_CMAPFILE_DIR is not None:
            cmapsflist = sorted(glob(os.path.join(USER_CMAPFILE_DIR, '*.rgb')))
            for cmap_file in cmapsflist:
                cname = os.path.basename(cmap_file).split('.rgb')[0]
                # start with the number will result illegal attribute
                if cname[0].isdigit() or cname.startswith('_'):
                    cname = 'C' + cname
                if '-' in cname:
                    cname = 'cmaps_' + cname.replace('-', '_')
                if '+' in cname:
                    cname = 'cmaps_' + cname.replace('+', '_')

                try:
                    cmap = get_cmap(cname)
                except:
                    cmap = Colormap(self._coltbl(cmap_file), name=cname)
                    register_cmap(name=cname, cmap=cmap)
                setattr(self, cname, cmap)

                cname = cname + '_r'
                try:
                    cmap = get_cmap(cname)
                except:
                    cmap = Colormap(self._coltbl(cmap_file)[::-1], name=cname)
                    register_cmap(name=cname, cmap=cmap)
                setattr(self, cname, cmap)