Review: a few fixes

[pyi2ncommon] / src / text_helpers.py

# The software in this package is distributed under the GNU General
# Public License version 2 (with a special exception described below).
#
# A copy of GNU General Public License (GPL) is included in this distribution,
# in the file COPYING.GPL.
#
# As a special exception, if other files instantiate templates or use macros
# or inline functions from this file, or you compile this file and link it
# with other works to produce a work based on this file, this file
# does not by itself cause the resulting work to be covered
# by the GNU General Public License.
#
# However the source code for this file must still be made available
# in accordance with section (3) of the GNU General Public License.
#
# This exception does not invalidate any other reasons why a work based
# on this file might be covered by the GNU General Public License.
#
# Copyright (c) 2016-2018 Intra2net AG <info@intra2net.com>

"""

SUMMARY
------------------------------------------------------
Functions for improving textual output.

Copyright: 2015 Intra2net AG


CONTENTS
------------------------------------------------------
This module has two parts. Part 1 includes:

- head_and_tail: shows the first few and last few elements of an iterable that
                 could potentially be pretty long
- size_str: textual representation of data size

Part2 contains functions for coloring text, a poor-man's version of other
modules like :py:mod:`colorclass` (which is now also available on Intra2net
systems)

Functions might cause trouble when combined, e.g.::

    bold('this is bold and ' + red('red') + ' and ' + green('green'))

will show the text "and green" not in bold. May have to try using specific
end-of-color or end-of-style escape sequences instead of 0 (reset-everything).


.. seealso:: http://stackoverflow.com/questions/287871/print-in-terminal-with-colors-using-python
.. seealso:: https://en.wikipedia.org/wiki/ANSI_escape_code


INTERFACE
------------------------------------------------------
"""

try:
    from builtins import print as _builtin_print
except ImportError:    # different name in py2
    from __builtin__ import print as _builtin_print

from functools import partial
from itertools import islice

from .type_helpers import isstr
from sys import stdout


def head_and_tail(iterable, n_head=20, n_tail=20, n_elems=None,
                  skip_elem="...(skipping {n_skipped} elements)..."):
    """
    Convenient way to shorten a possibly very long iterable before printing.

    Will not modify short iterables, but for longer lists/tuples/... will only
    yield first few, then a message how many were skipped and the last few

    The iterable does not even have to have a `len(..)` method if argument
    `n_elems` is given. Only needs a `next(..)` method. However, for very long
    iterables this will be faster if random element access is provided via []

    :param iterable: an iterable
    :type iterable: anything that can be iterated over
    :param int n_head: number of starting elements to yield (optional)
    :param int n_tail: number of ending elements to yield (optional)
    :param int n_elems: number of elements in iterable; give this to avoid a
                        call to `len(iterable)` (optional)
    :param skip_elem: element to replace bulge of skipped elements; yielded
                      once at most; None to not yield a skip replacement; if str
                      then it will be formatted; optional, defaults to string
                      with number of skipped elems
    :type skip_elem: anything you like
    :yields: `n_head+n_tail` elements from iterable plus the `skip_elem` (or
             less if iterable is shorter than this).

    .. seealso:: :py:func:`slice`, :py:func:`itertools.islice`
    """
    if n_elems is None:
        n_elems = len(iterable)

    # yield first n_head elems
    index = 0
    for elem in iterable:
        index += 1
        if index > n_head:
            break
        yield elem

    # yield skip element
    if n_elems > n_head + n_tail:
        if skip_elem is not None:
            if isstr(skip_elem):
                yield skip_elem.format(n_skipped=n_elems-n_head-n_tail)
            else:
                yield skip_elem
    elif n_elems <= n_head:
        return

    # yield end
    try:
        # try to access end directly
        for elem in iterable[n_elems-n_tail:]:
            yield elem
    except TypeError:
        # if this did not work, then need to iterate through whole thing
        # do this as in itertool recipe for consume():
        n_skip = n_elems - n_head - n_tail - 1
        next(islice(iterable, n_skip, n_skip), None)
        for elem in iterable:
            yield elem


def size_str(byte_number, is_diff=False):
    """
    Create a human-readable text representation of a file size.

    Rounds and shortens size to something easily human-readable like '1.5 GB'.

    :param int byte_number: Number of bytes to express as text
    :param bool is_diff: Set to True to include a '+' or '-' in output;
                         default: False
    :returns: textual representation of data
    :rtype: str
    """
    # constants
    units = '', 'k', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y'
    factor = 1024
    thresh_add_comma = 10.   # below this, return 1.2G, above this return 12G

    # prepare
    if byte_number < 0:
        sign_str = '-'
    elif is_diff:
        sign_str = '+'
    else:
        sign_str = ''
    curr_num = abs(float(byte_number))

    # loop
    for unit in units:
        if curr_num > factor:
            curr_num /= factor
            continue
        elif curr_num < thresh_add_comma and unit != 'B':   # e.g. 1.2G
            return '{2}{0:.1f} {1}B'.format(curr_num, unit, sign_str)
        else:   # e.g. 12G or 1B
            return '{2}{0:d} {1}B'.format(int(round(curr_num)), unit, sign_str)

    # have an impossible amount of data.  (>1024**4 GB)
    # undo last "/factor" and show thousand-separator
    return '{2}{0:,d} {1}B'.format(int(round(curr_num*factor)), units[-1],
                                 sign_str)


###############################################################################
# TEXT FORMATTING/COLORING
###############################################################################

COLOR_BLACK = 'black'
COLOR_RED = 'red'
COLOR_GREEN = 'green'
COLOR_YELLOW = 'yellow'
COLOR_BLUE = 'blue'
COLOR_MAGENTA = 'magenta'
COLOR_CYAN = 'cyan'
COLOR_WHITE = 'white'

STYLE_NORMAL = 0
STYLE_BOLD = 1
STYLE_UNDERLINE = 4
STYLE_BLINK = 5
STYLE_REVERSE = 7


_COLOR_TO_CODE = dict(zip([COLOR_BLACK, COLOR_RED, COLOR_GREEN, COLOR_YELLOW, \
                           COLOR_BLUE, COLOR_MAGENTA, COLOR_CYAN, COLOR_WHITE],
                          range(8)))

_ANSI_ESCAPE_SURROUND = '\x1b[{}m{}\x1b[0m'

# only color output if we are writing output to a terminal (not a file or so)
try:
    _STDOUT_IS_TTY = stdout.isatty()
except Exception:
    # stdout might be some wrapper around the real thing to capture output
    _STDOUT_IS_TTY = False


def colored(text, foreground=None, background=None, style=None):
    """ return text with given foreground/background ANSI color escape seqence

    :param str text: text to color
    :param str style: one of the style constants above
    :param str foreground: one of the color constants to use for text color
                           or None to leave as-is
    :param str foreground: one of the color constants to use for background
                           or None to leave as-is
    :param style: single STYLE constant or iterable of those
                  or None to leave as-is
    :returns: text surrounded in ansi escape sequences
    """

    if foreground is None and background is None and style is None:
        return text

    prefixes = []
    if foreground:
        prefixes.append(str(30 + _COLOR_TO_CODE[foreground]))
    if background:
        prefixes.append(str(40 + _COLOR_TO_CODE[background]))
    if style is None:
        pass
    elif isinstance(style, int):
        prefixes.append(str(style))
    else:   # assume iterable of ints
        prefixes.extend(str(style_item) for style_item in style)

    return _ANSI_ESCAPE_SURROUND.format(';'.join(prefixes), text)


def print(text, *args, **kwargs):           # pylint: disable=redefined-builtin
    """ like the builtin print function but also accepts color args

    If any arg of :py:func:`colored` is given in `kwargs`, will run text with
    color-args through that function. Runs built-in :py:func:`print`
    function with result and other args.

    ...todo:: color all args, not just first
    """
    foreground = None
    background = None
    style = None

    # remove color info from kwargs
    try:
        foreground = kwargs['foreground']
        del kwargs['foreground']
    except KeyError:
        pass

    try:
        background = kwargs['background']
        del kwargs['background']
    except KeyError:
        pass

    try:
        style = kwargs['style']
        del kwargs['style']
    except KeyError:
        pass

    if _STDOUT_IS_TTY:
        text_c = colored(text, foreground, background, style)
    else:
        text_c = text

    _builtin_print(text_c, *args, **kwargs)


black   = partial(print, foreground=COLOR_BLACK)
red     = partial(print, foreground=COLOR_RED)
green   = partial(print, foreground=COLOR_GREEN)
yellow  = partial(print, foreground=COLOR_YELLOW)
blue    = partial(print, foreground=COLOR_BLUE)
magenta = partial(print, foreground=COLOR_MAGENTA)
cyan    = partial(print, foreground=COLOR_CYAN)
white   = partial(print, foreground=COLOR_WHITE)

normal    = partial(print, style=STYLE_NORMAL)
bold      = partial(print, style=STYLE_BOLD)
underline = partial(print, style=STYLE_UNDERLINE)
blink     = partial(print, style=STYLE_BLINK)
reverse   = partial(print, style=STYLE_REVERSE)