Jinja2 Template Tester

j2.py

"""Jinja2 Template Tester."""

import sys
import re
import os.path
import argparse
import logging
import inspect

import jinja2


class DefaultHelpParser(argparse.ArgumentParser):
    """Prints help on argparse error."""

    def error(self, message):
        """Print help message on error.

        Args:
            message: Help message to display.
        """
        script = os.path.basename(__file__)
        sys.stderr.write('{}: error: {}\n'.format(script, message))
        self.print_help(sys.stderr)
        sys.exit(2)


class Stub:
    """Stub class to swallow Jinja2 exceptions."""

    def __init__(self):
        pass

    def __trunc__(self):  # pylint: disable=no-self-use
        return 0

    def __iter__(self):
        return iter(())

    def next(self):  # pylint: disable=no-self-use,missing-docstring
        raise StopIteration

    def __repr__(self):
        return ''

    def __call__(self):
        return self

    __str__ = __repr__


class SilentUndefined(jinja2.Undefined):
    # pylint: disable=too-few-public-methods
    """Make Jinja2 undefined errors silent.

    Based on https://stackoverflow.com/q/6190348/5764537.
    """
    def _fail_with_undefined_error(self, *_args, **_kwargs):
        """Inspect stack to find undefined variables in context."""
        for func in inspect.stack():
            f_locals = func[0].f_locals
            if any(key[:2] == 'l_' for key in f_locals.keys()):
                undefs = [name for name in map(
                    lambda k, f_l=f_locals: k[2:] if k[:2] == 'l_'
                    and not isinstance(f_l[k], str) else None,
                    f_locals.keys()) if name is not None]
                logging.debug('Undefined context variables: %s',
                              ', '.join(map(lambda u: u[2:], undefs)))
                break
        return Stub()
    __add__ = __radd__ = __mul__ = __rmul__ = __div__ = __rdiv__ = \
        __truediv__ = __rtruediv__ = __floordiv__ = __rfloordiv__ = \
        __mod__ = __rmod__ = __pos__ = __neg__ = __call__ = \
        __getitem__ = __lt__ = __le__ = __gt__ = __ge__ = \
        __complex__ = __pow__ = __rpow__ = \
        _fail_with_undefined_error
    __int__ = __index__ = lambda self, *args, **kwargs: int()
    __float__ = lambda self, *args, **kwargs: float()


def is_binary(path):
    """Return true if the given file is binary.

    Based on https://stackoverflow.com/q/898669/5764537.

    Args:
        path: Path to file to check.

    Returns:
        True if the file is binary, False otherwise.
    """
    file = open(path, 'rb')
    try:
        chunk_size = 1024
        while 1:
            chunk = file.read(chunk_size)
            if b'\0' in chunk:
                return True
            if len(chunk) < chunk_size:
                break
    finally:
        file.close()
    return False


def is_int(literal):
    """Return True if the given string is an integer.

    Based on: https://stackoverflow.com/q/1265665/5764537.

    Args:
        literal: The input string to check.

    Returns:
        True if input string is an integer, False otherwise.
    """
    try:
        int(literal)
        return True
    except ValueError:
        return False


def set_value(dictionary, key, value):
    """Expand dot notation in key and set the value in dictionary.

    Args:
        dictionary: The dict to update.
        key: Key to set with dot notation.
        value: Value to set.
    """
    # autovivify path defined with key
    for element in key.split('.')[:-1]:
        dictionary = dictionary.setdefault(element, dict())
    # convert value to int if possible
    if is_int(value):
        value = int(value)
    # modify dict in place
    dictionary[key.split('.')[-1]] = value


def create_jinja2_context(context_args):
    """Create Jinja2 context from the command line arguments.

    Args:
        context_args: Context arguments in the format from the command line.
            For example foo:bar or foo.bar.baz:qux.

    Returns:
        A dictionary that can be used as a Jinja2 context.
    """
    context = dict()
    for arg in context_args:
        try:
            key, value = arg.split('=')
        except ValueError:
            logging.error(
                'Context key/value pair "%s" not in either of the following '
                'formats: foo:bar, foo.bar.baz:qux', arg)
            continue
        set_value(context, key, value)
    return context


def find_templates(path):
    """Find all the templates based on the path.

    Args:
        path: Path to search for templates. Can be a single filename
            or a directory.

    Returns:
        List of paths to potential Jinja2 templates.
    """
    files = []
    if os.path.isfile(path):
        # path is a filename
        files = [path]
    elif os.path.isdir(path):
        # path is a directory, find potential templates inside
        files = [os.path.join(root, name)
                 for root, dirs, names in os.walk(path)
                 for name in names]
    else:
        logging.critical('"%s" is not a directory or a file', path)
        sys.exit(1)
    # filter binary files
    files = [name for name in map(
        lambda file: file if not is_binary(file) else None,
        files) if name is not None]
    return files


@jinja2.environmentfilter
def jinja2_filter_stub(_, value):
    """Jinja2 filter stub to swallow missing filter errors."""
    return value


def parse_args():
    """Parse arguments."""
    parser = DefaultHelpParser(description='Test Python Jinja2 templates with '
                                           'high tolerance for errors.')
    parser.add_argument('template', metavar='template', type=str,
                        help='template file or directory')
    parser.add_argument('context', metavar='key=value', type=str, nargs='*',
                        help='context key/value pair, eg. foo:bar or '
                        'foo.bar.baz:qux')
    parser.add_argument('--debug', action='store_true', help='debug mode')
    return parser.parse_args()


def main():
    """Jinja2 Template Tester."""
    opts = parse_args()
    logging.basicConfig(format='%(levelname)s:%(message)s',
                        level=logging.DEBUG if opts.debug else logging.INFO)
    context = create_jinja2_context(opts.context)
    logging.debug('Using context %s', context)
    filenames = find_templates(opts.template)
    logging.debug('Using templates %s', ', '.join(filenames))
    ignore_filters = []
    for name in filenames:
        with open(name) as file:
            template_str = file.read()
            env = jinja2.Environment(loader=jinja2.BaseLoader,
                                     undefined=SilentUndefined)
            while True:
                try:
                    template = env.from_string(template_str)
                except jinja2.exceptions.TemplateAssertionError as exc:
                    logging.debug(exc.message)
                    try:
                        # check for missing filter errors
                        filter_name = re.search(
                            "no filter named '(.+)'", exc.message).group(1)
                        # add missing filter as a stub
                        ignore_filters.append(filter_name)
                        logging.debug('Adding filter stub for %s', filter_name)
                        env.filters.update({
                            filter: jinja2_filter_stub
                            for filter in ignore_filters})
                    except AttributeError as exc:
                        logging.debug(exc.message)
                        continue
                else:
                    break
            print('{}\n{}'.format(
                name + ':', template.render(context)))


if __name__ == '__main__':
    main()