Source code for cli2.display

"""
Generic pretty display utils.

.. envvar:: FORCE_COLOR

    By default, we will not color strings in non-interactive ttys, but you can
    force it with :envvar:`FORCE_COLOR`, ie. gitlab-ci etc

.. envvar:: CLI2_PYGMENTS_STYLE

    Pygments style to use when highlighting code, monokai by default.
"""
import difflib
import json
import os
import sys
import yaml

_print = print


def color_enabled():
    if 'FORCE_COLOR' in os.environ:
        return bool(os.getenv('FORCE_COLOR', ''))
    return sys.stdout.isatty()




def highlight(string, lexer):
    """
    Use pygments to render a string with a lexer.

    :param string: String to render
    :param lexer: Lexer name, Yaml, Diff, etc
    """
    if not color_enabled():
        return string

    try:
        import pygments
        import pygments.lexers
        import pygments.formatters
    except ImportError:
        return string

    formatter = pygments.formatters.Terminal256Formatter(
        style=os.getenv('CLI2_PYGMENTS_STYLE', 'monokai'),
    )
    lexer = getattr(pygments.lexers, lexer + 'Lexer')()
    return pygments.highlight(string, lexer, formatter).rstrip()



def yaml_dump(data):
    if isinstance(data, dict):
        # ensure that objects inheriting from dict render nicely
        data = dict(data)
    return yaml.dump(data, indent=4, width=float('inf'))




def yaml_highlight(yaml_string):
    """ Return highlighted YAML """
    return highlight(yaml_string, 'Yaml')





def render(arg, highlight=True):
    """
    Try to render arg as yaml.

    If the arg has a ``.json()`` method, it'll be called.
    If it is parseable as JSON then it'l be parsed as such.
    Then, it'll be dumped as colored YAML.

    Set the env var `FORCE_COLOR` to anything to force into printing colors
    even if terminal is non-interactive (ie. gitlab-ci)

    .. code-block:: python

        # pretty render some_object
        print(cli2.render(some_object))
    """
    try:  # deal with response objects
        arg = arg.json()
    except:  # noqa
        pass

    try:  # is this json?
        arg = json.loads(arg)
    except:  # noqa
        pass

    # does this wants to show specific data to cli2?
    try:
        arg = arg.cli2_display
    except AttributeError:
        pass

    string = arg if isinstance(arg, str) else yaml_dump(arg)
    if not highlight:
        return string
    return yaml_highlight(string)





def print(*args, **kwargs):
    """
    Try to print the :py:func:`render`'ed args, pass the kwargs to actual print
    method.

    .. code-block:: python

        # pretty print some_object
        cli2.print(some_object)
    """
    for arg in args:
        _print(render(arg), **kwargs)





def diff_highlight(diff):
    """
    Return highlighted unified diff.
    """
    output = '\n'.join([line.rstrip() for line in diff if line.strip()])
    return highlight(output, 'Diff')





def diff(diff, **kwargs):
    """
    Pretty-print a diff generated by Python's standard difflib.unified_diff
    method.

    .. code-block:: python

        # pretty print a diff
        cli2.diff(difflib.unified_diff(old, new))
    """
    _print(diff_highlight(diff), **kwargs)





def diff_data(before, after, before_label='before', after_label='after'):
    """
    YAML Dump before and after objects and return the unified diff, printable
    with :py:func:`diff`.

    :param before: First object to compare
    :param after: Second object to compare
    :param before_label: Name of the first object to display in diff
    :param after_label: Name of the second object to display in diff
    """
    return difflib.unified_diff(
        yaml.dump(before).splitlines(),
        yaml.dump(after).splitlines(),
        before_label,
        after_label,
    )