pygments_cache.py

"""A fast, drop-in replacement for pygments ``get_*()`` and ``guess_*()`` funtions.

The following pygments API functions are currently supplied here::

    from pygments_cache import get_lexer_for_filename, guess_lexer_for_filename
    from pygments_cache import get_formatter_for_filename, get_formatter_by_name
    from pygments_cache import get_style_by_name, get_all_styles
    from pygments_cache import get_filter_by_name

The cache itself is stored at the location given by the ``$PYGMENTS_CACHE_FILE``
environment variable, or by default at ``~/.local/share/pygments-cache/cache.py``.
The cache file is created on first use, if it does not already exist.


"""
import os
import importlib


# Global storage variables
__version__ = '0.1.3'
CACHE = None
CUSTOM_STYLES = {}
DEBUG = False


def _print_duplicate_message(duplicates):
    import sys
    for filename, vals in sorted(duplicates.items()):
        msg = 'for {0} ambiquity between:\n  '.format(filename)
        vals = [m + ':' + c for m, c in vals]
        msg += '\n  '.join(sorted(vals))
        print(msg, file=sys.stderr)


def _discover_lexers():
    import inspect
    from pygments.lexers import get_all_lexers, find_lexer_class
    # maps file extension (and names) to (module, classname) tuples
    default_exts = {
        # C / C++
        '.h': ('pygments.lexers.c_cpp', 'CLexer'),
        '.hh': ('pygments.lexers.c_cpp', 'CppLexer'),
        '.cp': ('pygments.lexers.c_cpp', 'CppLexer'),
        # python
        '.py': ('pygments.lexers.python', 'Python3Lexer'),
        '.pyw': ('pygments.lexers.python', 'Python3Lexer'),
        '.sc': ('pygments.lexers.python', 'Python3Lexer'),
        '.tac': ('pygments.lexers.python', 'Python3Lexer'),
        'SConstruct': ('pygments.lexers.python', 'Python3Lexer'),
        'SConscript': ('pygments.lexers.python', 'Python3Lexer'),
        '.sage': ('pygments.lexers.python', 'Python3Lexer'),
        '.pytb': ('pygments.lexers.python', 'Python3TracebackLexer'),
        # perl
        '.t': ('pygments.lexers.perl', 'Perl6Lexer'),
        '.pl': ('pygments.lexers.perl', 'Perl6Lexer'),
        '.pm': ('pygments.lexers.perl', 'Perl6Lexer'),
        # asm
        '.s': ('pygments.lexers.asm', 'GasLexer'),
        '.S': ('pygments.lexers.asm', 'GasLexer'),
        '.asm': ('pygments.lexers.asm', 'NasmLexer'),
        '.ASM': ('pygments.lexers.asm', 'NasmLexer'),
        # Antlr
        '.g': ('pygments.lexers.parsers', 'AntlrCppLexer'),
        '.G': ('pygments.lexers.parsers', 'AntlrCppLexer'),
        # XML
        '.xml': ('pygments.lexers.html', 'XmlLexer'),
        '.xsl': ('pygments.lexers.html', 'XsltLexer'),
        '.xslt': ('pygments.lexers.html', 'XsltLexer'),
        # ASP
        '.axd': ('pygments.lexers.dotnet', 'CSharpAspxLexer'),
        '.asax': ('pygments.lexers.dotnet', 'CSharpAspxLexer'),
        '.ascx': ('pygments.lexers.dotnet', 'CSharpAspxLexer'),
        '.ashx': ('pygments.lexers.dotnet', 'CSharpAspxLexer'),
        '.asmx': ('pygments.lexers.dotnet', 'CSharpAspxLexer'),
        '.aspx': ('pygments.lexers.dotnet', 'CSharpAspxLexer'),
        # misc
        '.b': ('pygments.lexers.esoteric', 'BrainfuckLexer'),
        '.j': ('pygments.lexers.jvm', 'JasminLexer'),
        '.m': ('pygments.lexers.matlab', 'MatlabLexer'),
        '.n': ('pygments.lexers.dotnet', 'NemerleLexer'),
        '.p': ('pygments.lexers.pawn', 'PawnLexer'),
        '.v': ('pygments.lexers.theorem', 'CoqLexer'),
        '.as': ('pygments.lexers.actionscript', 'ActionScript3Lexer'),
        '.fs': ('pygments.lexers.forth', 'ForthLexer'),
        '.hy': ('pygments.lexers.lisp', 'HyLexer'),
        '.ts': ('pygments.lexers.javascript', 'TypeScriptLexer'),
        '.rl': ('pygments.lexers.parsers', 'RagelCppLexer'),
        '.bas': ('pygments.lexers.basic', 'QBasicLexer'),
        '.bug': ('pygments.lexers.modeling', 'BugsLexer'),
        '.ecl': ('pygments.lexers.ecl', 'ECLLexer'),
        '.inc': ('pygments.lexers.php', 'PhpLexer'),
        '.inf': ('pygments.lexers.configs', 'IniLexer'),
        '.pro': ('pygments.lexers.prolog', 'PrologLexer'),
        '.sql': ('pygments.lexers.sql', 'SqlLexer'),
        '.txt': ('pygments.lexers.special', 'TextLexer'),
        '.html': ('pygments.lexers.html', 'HtmlLexer'),
        }
    exts = {}
    lexers = {'exts': exts}
    if DEBUG:
        from collections import defaultdict
        duplicates = defaultdict(set)
    for longname, aliases, filenames, mimetypes in get_all_lexers():
        cls = find_lexer_class(longname)
        mod = inspect.getmodule(cls)
        val = (mod.__name__, cls.__name__)
        for filename in filenames:
            if filename.startswith('*.'):
                filename = filename[1:]
            if '*' in filename:
                continue
            if (DEBUG and filename in exts and exts[filename] != val
                    and filename not in default_exts):
                duplicates[filename].add(val)
                duplicates[filename].add(exts[filename])
            exts[filename] = val
    # remove some ambiquity
    exts.update(default_exts)
    # print duplicate message
    if DEBUG:
        _print_duplicate_message(duplicates)
    return lexers


def _discover_formatters():
    import inspect
    from pygments.formatters import get_all_formatters
    # maps file extension (and names) to (module, classname) tuples
    default_exts = {}
    exts = {}
    # maps formatter 'name' (not the class name) and alias to (module, classname) tuples
    default_names = {}
    names = {}
    formatters = {'exts': exts, 'names': names}
    if DEBUG:
        from collections import defaultdict
        duplicates = defaultdict(set)
    for cls in get_all_formatters():
        mod = inspect.getmodule(cls)
        val = (mod.__name__, cls.__name__)
        # add extentions
        for filename in cls.filenames:
            if filename.startswith('*.'):
                filename = filename[1:]
            if '*' in filename:
                continue
            if (DEBUG and filename in exts and exts[filename] != val
                    and filename not in default_exts):
                duplicates[filename].add(val)
                duplicates[filename].add(exts[filename])
            exts[filename] = val
        # add names and aliases
        names[cls.name] = val
        for alias in cls.aliases:
            if (DEBUG and alias in names and names[alias] != val
                    and alias not in default_names):
                duplicates[alias].add(val)
                duplicates[alias].add(names[alias])
            names[alias] = val
    # remove some ambiquity
    exts.update(default_exts)
    names.update(default_names)
    # print dumplicate message
    if DEBUG:
        _print_duplicate_message(duplicates)
    return formatters


def _discover_styles():
    import inspect
    from pygments.styles import get_all_styles, get_style_by_name
    # maps style 'name' (not the class name) and aliases to (module, classname) tuples
    default_names = {}
    names = {}
    styles = {'names': names}
    if DEBUG:
        from collections import defaultdict
        duplicates = defaultdict(set)
    for name in get_all_styles():
        cls = get_style_by_name(name)
        mod = inspect.getmodule(cls)
        val = (mod.__name__, cls.__name__)
        if (DEBUG and name in names and names[name] != val
                and name not in default_names):
            duplicates[name].add(val)
            duplicates[name].add(names[name])
        names[name] = val
    # remove some ambiquity
    names.update(default_names)
    # print dumplicate message
    if DEBUG:
        _print_duplicate_message(duplicates)
    return styles


def _discover_filters():
    import inspect
    from pygments.filters import get_all_filters, get_filter_by_name
    # maps filter 'name' (not the class name) to (module, classname) tuples
    default_names = {}
    names = {}
    filters = {'names': names}
    if DEBUG:
        from collections import defaultdict
        duplicates = defaultdict(set)
    for name in get_all_filters():
        filter = get_filter_by_name(name)
        cls = type(filter)
        mod = inspect.getmodule(cls)
        val = (mod.__name__, cls.__name__)
        if (DEBUG and name in names and names[name] != val
                and name not in default_names):
            duplicates[name].add(val)
            duplicates[name].add(names[name])
        names[name] = val
    # remove some ambiquity
    names.update(default_names)
    # print dumplicate message
    if DEBUG:
        _print_duplicate_message(duplicates)
    return filters


def build_cache():
    """Does the hard work of building a cache from nothing."""
    cache = {}
    cache['lexers'] = _discover_lexers()
    cache['formatters'] = _discover_formatters()
    cache['styles'] = _discover_styles()
    cache['filters'] = _discover_filters()
    return cache


def cache_filename():
    """Gets the name of the cache file to use."""
    # Configuration variables read from the environment
    if 'PYGMENTS_CACHE_FILE' in os.environ:
        return os.environ['PYGMENTS_CACHE_FILE']
    else:
        return os.path.join(os.environ.get('XDG_DATA_HOME',
                                           os.path.join(os.path.expanduser('~'),
                                                        '.local', 'share')),
                            'pygments-cache',
                            'cache.py'
                            )


def add_custom_style(name, style):
    """Register custom style to be able to retrieve it by ``get_style_by_name``.

    Parameters
    ----------
    name : str
        Style name.
    style : pygments.Style
        Custom style to add.
    """
    CUSTOM_STYLES[name] = style


def load(filename):
    """Loads the cache from a filename."""
    global CACHE
    with open(filename) as f:
        s = f.read()
    ctx = globals()
    CACHE = eval(s, ctx, ctx)
    return CACHE


def write_cache(filename):
    """Writes the current cache to the file"""
    from pprint import pformat
    d = os.path.dirname(filename)
    os.makedirs(d, exist_ok=True)
    s = pformat(CACHE)
    with open(filename, 'w') as f:
        f.write(s)


def load_or_build():
    """Loads the cache from disk. If the cache does not exist,
    this will build and write it out.
    """
    global CACHE
    fname = cache_filename()
    if os.path.exists(fname):
        load(fname)
    else:
        import sys
        if DEBUG:
            print('pygments cache not found, building...', file=sys.stderr)
        CACHE = build_cache()
        if DEBUG:
            print('...writing cache to ' + fname, file=sys.stderr)
        write_cache(fname)


#
# pygments interface
#

def get_lexer_for_filename(filename, text='', **options):
    """Gets a lexer from a filename (usually via the filename extension).
    This mimics the behavior of ``pygments.lexers.get_lexer_for_filename()``
    and ``pygments.lexers.guess_lexer_for_filename()``.
    """
    if CACHE is None:
        load_or_build()
    exts = CACHE['lexers']['exts']
    fname = os.path.basename(filename)
    key = fname if fname in exts else os.path.splitext(fname)[1]
    if key in exts:
        modname, clsname = exts[key]
        mod = importlib.import_module(modname)
        cls = getattr(mod, clsname)
        lexer = cls(**options)
    else:
        # couldn't find lexer in cache, fallback to the hard way
        import inspect
        from pygments.lexers import guess_lexer_for_filename
        lexer = guess_lexer_for_filename(filename, text, **options)
        # add this filename to the cache for future use
        cls = type(lexer)
        mod = inspect.getmodule(cls)
        exts[fname] = (mod.__name__, cls.__name__)
        write_cache(cache_filename())
    return lexer


guess_lexer_for_filename = get_lexer_for_filename


def get_formatter_for_filename(fn, **options):
    """Gets a formatter instance from a filename (usually via the filename
    extension). This mimics the behavior of
    ``pygments.formatters.get_formatter_for_filename()``.
    """
    if CACHE is None:
        load_or_build()
    exts = CACHE['formatters']['exts']
    fname = os.path.basename(fn)
    key = fname if fname in exts else os.path.splitext(fname)[1]
    if key in exts:
        modname, clsname = exts[key]
        mod = importlib.import_module(modname)
        cls = getattr(mod, clsname)
        formatter = cls(**options)
    else:
        # couldn't find formatter in cache, fallback to the hard way
        import inspect
        from pygments.formatters import get_formatter_for_filename
        formatter = get_formatter_for_filename(fn, **options)
        # add this filename to the cache for future use
        cls = type(formatter)
        mod = inspect.getmodule(cls)
        exts[fname] = (mod.__name__, cls.__name__)
        write_cache(cache_filename())
    return formatter


def get_formatter_by_name(alias, **options):
    """Gets a formatter instance from its name or alias.
    This mimics the behavior of ``pygments.formatters.get_formatter_by_name()``.
    """
    if CACHE is None:
        load_or_build()
    names = CACHE['formatters']['names']
    if alias in names:
        modname, clsname = names[alias]
        mod = importlib.import_module(modname)
        cls = getattr(mod, clsname)
        formatter = cls(**options)
    else:
        # couldn't find formatter in cache, fallback to the hard way
        import inspect
        from pygments.formatters import get_formatter_by_name
        formatter = get_formatter_by_name(alias, **options)
        # add this filename to the cache for future use
        cls = type(formatter)
        mod = inspect.getmodule(cls)
        names[alias] = (mod.__name__, cls.__name__)
        write_cache(cache_filename())
    return formatter


def get_style_by_name(name):
    """Gets a style class from its name or alias.
    This mimics the behavior of ``pygments.styles.get_style_by_name()``.
    """
    if CACHE is None:
        load_or_build()
    names = CACHE['styles']['names']
    if name in names:
        modname, clsname = names[name]
        mod = importlib.import_module(modname)
        style = getattr(mod, clsname)
    elif name in CUSTOM_STYLES:
        style = CUSTOM_STYLES[name]
    else:
        # couldn't find style in cache, fallback to the hard way
        import inspect
        from pygments.styles import get_style_by_name
        style = get_style_by_name(name)
        # add this style to the cache for future use
        mod = inspect.getmodule(style)
        names[name] = (mod.__name__, style.__name__)
        write_cache(cache_filename())
    return style


def get_all_styles():
    """Iterable through all known style names.
    This mimics the behavior of ``pygments.styles.get_all_styles``.
    """
    if CACHE is None:
        load_or_build()
    yield from CACHE['styles']['names']
    yield from CUSTOM_STYLES


def get_filter_by_name(filtername, **options):
    """Gets a filter instance from its name. This mimics the behavior of
    ``pygments.filters.get_filter_by_name()``.
    """
    if CACHE is None:
        load_or_build()
    names = CACHE['filters']['names']
    if filtername in names:
        modname, clsname = names[filtername]
        mod = importlib.import_module(modname)
        cls = getattr(mod, clsname)
        filter = cls(**options)
    else:
        # couldn't find style in cache, fallback to the hard way
        import inspect
        from pygments.filters import get_filter_by_name
        filter = get_filter_by_name(filtername, **options)
        # add this filter to the cache for future use
        cls = type(filter)
        mod = inspect.getmodule(cls)
        names[filtername] = (mod.__name__, cls.__name__)
        write_cache(cache_filename())
    return filter