developer.intra2net.com Git - pyi2ncommon/blob - src/text_helpers.py

   1 # The software in this package is distributed under the GNU General
   2 # Public License version 2 (with a special exception described below).
   3 #
   4 # A copy of GNU General Public License (GPL) is included in this distribution,
   5 # in the file COPYING.GPL.
   6 #
   7 # As a special exception, if other files instantiate templates or use macros
   8 # or inline functions from this file, or you compile this file and link it
   9 # with other works to produce a work based on this file, this file
  10 # does not by itself cause the resulting work to be covered
  11 # by the GNU General Public License.
  12 #
  13 # However the source code for this file must still be made available
  14 # in accordance with section (3) of the GNU General Public License.
  15 #
  16 # This exception does not invalidate any other reasons why a work based
  17 # on this file might be covered by the GNU General Public License.
  18 #
  19 # Copyright (c) 2016-2018 Intra2net AG <info@intra2net.com>
  20
  21 """
  22
  23 SUMMARY
  24 ------------------------------------------------------
  25 Functions for improving textual output.
  26
  27 Copyright: 2015 Intra2net AG
  28
  29
  30 CONTENTS
  31 ------------------------------------------------------
  32 This module has two parts. Part 1 includes:
  33
  34 - head_and_tail: shows the first few and last few elements of an iterable that
  35                  could potentially be pretty long
  36 - size_str: textual representation of data size
  37
  38 Part2 contains functions for coloring text, a poor-man's version of other
  39 modules like :py:mod:`colorclass` (which is now also available on Intra2net
  40 systems)
  41
  42 Functions might cause trouble when combined, e.g.::
  43
  44     bold('this is bold and ' + red('red') + ' and ' + green('green'))
  45
  46 will show the text "and green" not in bold. May have to try using specific
  47 end-of-color or end-of-style escape sequences instead of 0 (reset-everything).
  48
  49
  50 .. seealso:: http://stackoverflow.com/questions/287871/print-in-terminal-with-colors-using-python
  51 .. seealso:: https://en.wikipedia.org/wiki/ANSI_escape_code
  52 .. seealso:: :py:mod:`textwrap`
  53
  54
  55 INTERFACE
  56 ------------------------------------------------------
  57 """
  58
  59 try:
  60     from builtins import print as _builtin_print
  61 except ImportError:    # different name in py2
  62     from __builtin__ import print as _builtin_print
  63
  64 from functools import partial
  65 from itertools import islice
  66
  67 from .type_helpers import isstr
  68 from sys import stdout
  69
  70
  71 def head_and_tail(iterable, n_head=20, n_tail=20, n_elems=None,
  72                   skip_elem="...(skipping {n_skipped} elements)..."):
  73     """
  74     Convenient way to shorten a possibly very long iterable before printing.
  75
  76     Will not modify short iterables, but for longer lists/tuples/... will only
  77     yield first few, then a message how many were skipped and the last few
  78
  79     The iterable does not even have to have a `len(..)` method if argument
  80     `n_elems` is given. Only needs a `next(..)` method. However, for very long
  81     iterables this will be faster if random element access is provided via []
  82
  83     :param iterable: an iterable
  84     :type iterable: anything that can be iterated over
  85     :param int n_head: number of starting elements to yield (optional)
  86     :param int n_tail: number of ending elements to yield (optional)
  87     :param int n_elems: number of elements in iterable; give this to avoid a
  88                         call to `len(iterable)` (optional)
  89     :param skip_elem: element to replace bulge of skipped elements; yielded
  90                       once at most; None to not yield a skip replacement; if str
  91                       then it will be formatted; optional, defaults to string
  92                       with number of skipped elems
  93     :type skip_elem: anything you like
  94     :yields: `n_head+n_tail` elements from iterable plus the `skip_elem` (or
  95              less if iterable is shorter than this).
  96
  97     .. seealso:: :py:func:`slice`, :py:func:`itertools.islice`, :py:func:`textwrap.shorten`
  98
  99     """
 100     if n_elems is None:
 101         n_elems = len(iterable)
 102
 103     # yield first n_head elems
 104     index = 0
 105     for elem in iterable:
 106         index += 1
 107         if index > n_head:
 108             break
 109         yield elem
 110
 111     # yield skip element
 112     if n_elems > n_head + n_tail:
 113         if skip_elem is not None:
 114             if isstr(skip_elem):
 115                 yield skip_elem.format(n_skipped=n_elems-n_head-n_tail)
 116             else:
 117                 yield skip_elem
 118     elif n_elems <= n_head:
 119         return
 120
 121     # yield end
 122     try:
 123         # try to access end directly
 124         for elem in iterable[n_elems-n_tail:]:
 125             yield elem
 126     except TypeError:
 127         # if this did not work, then need to iterate through whole thing
 128         # do this as in itertool recipe for consume():
 129         n_skip = n_elems - n_head - n_tail - 1
 130         next(islice(iterable, n_skip, n_skip), None)
 131         for elem in iterable:
 132             yield elem
 133
 134
 135 def size_str(byte_number, is_diff=False):
 136     """
 137     Create a human-readable text representation of a file size.
 138
 139     Rounds and shortens size to something easily human-readable like '1.5 GB'.
 140
 141     :param int byte_number: Number of bytes to express as text
 142     :param bool is_diff: Set to True to include a '+' or '-' in output;
 143                          default: False
 144     :returns: textual representation of data
 145     :rtype: str
 146     """
 147     # constants
 148     units = '', 'k', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y'
 149     factor = 1024
 150     thresh_add_comma = 10.   # below this, return 1.2G, above this return 12G
 151
 152     # prepare
 153     if byte_number < 0:
 154         sign_str = '-'
 155     elif is_diff:
 156         sign_str = '+'
 157     else:
 158         sign_str = ''
 159     curr_num = abs(float(byte_number))
 160
 161     # loop
 162     for unit in units:
 163         if curr_num > factor:
 164             curr_num /= factor
 165             continue
 166         elif curr_num < thresh_add_comma and unit != 'B':   # e.g. 1.2G
 167             return '{2}{0:.1f} {1}B'.format(curr_num, unit, sign_str)
 168         else:   # e.g. 12G or 1B
 169             return '{2}{0:d} {1}B'.format(int(round(curr_num)), unit, sign_str)
 170
 171     # have an impossible amount of data.  (>1024**4 GB)
 172     # undo last "/factor" and show thousand-separator
 173     return '{2}{0:,d} {1}B'.format(int(round(curr_num*factor)), units[-1],
 174                                  sign_str)
 175
 176
 177 ###############################################################################
 178 # TEXT FORMATTING/COLORING
 179 ###############################################################################
 180
 181 COLOR_BLACK = 'black'
 182 COLOR_RED = 'red'
 183 COLOR_GREEN = 'green'
 184 COLOR_YELLOW = 'yellow'
 185 COLOR_BLUE = 'blue'
 186 COLOR_MAGENTA = 'magenta'
 187 COLOR_CYAN = 'cyan'
 188 COLOR_WHITE = 'white'
 189
 190 STYLE_NORMAL = 0
 191 STYLE_BOLD = 1
 192 STYLE_UNDERLINE = 4
 193 STYLE_BLINK = 5
 194 STYLE_REVERSE = 7
 195
 196
 197 _COLOR_TO_CODE = dict(zip([COLOR_BLACK, COLOR_RED, COLOR_GREEN, COLOR_YELLOW, \
 198                            COLOR_BLUE, COLOR_MAGENTA, COLOR_CYAN, COLOR_WHITE],
 199                           range(8)))
 200
 201 _ANSI_ESCAPE_SURROUND = '\x1b[{}m{}\x1b[0m'
 202
 203 # only color output if we are writing output to a terminal (not a file or so)
 204 try:
 205     _STDOUT_IS_TTY = stdout.isatty()
 206 except Exception:
 207     # stdout might be some wrapper around the real thing to capture output
 208     _STDOUT_IS_TTY = False
 209
 210
 211 def colored(text, foreground=None, background=None, style=None):
 212     """ return text with given foreground/background ANSI color escape seqence
 213
 214     :param str text: text to color
 215     :param str style: one of the style constants above
 216     :param str foreground: one of the color constants to use for text color
 217                            or None to leave as-is
 218     :param str foreground: one of the color constants to use for background
 219                            or None to leave as-is
 220     :param style: single STYLE constant or iterable of those
 221                   or None to leave as-is
 222     :returns: text surrounded in ansi escape sequences
 223     """
 224
 225     if foreground is None and background is None and style is None:
 226         return text
 227
 228     prefixes = []
 229     if foreground:
 230         prefixes.append(str(30 + _COLOR_TO_CODE[foreground]))
 231     if background:
 232         prefixes.append(str(40 + _COLOR_TO_CODE[background]))
 233     if style is None:
 234         pass
 235     elif isinstance(style, int):
 236         prefixes.append(str(style))
 237     else:   # assume iterable of ints
 238         prefixes.extend(str(style_item) for style_item in style)
 239
 240     return _ANSI_ESCAPE_SURROUND.format(';'.join(prefixes), text)
 241
 242
 243 def print(text, *args, **kwargs):           # pylint: disable=redefined-builtin
 244     """ like the builtin print function but also accepts color args
 245
 246     If any arg of :py:func:`colored` is given in `kwargs`, will run text with
 247     color-args through that function. Runs built-in :py:func:`print`
 248     function with result and other args.
 249
 250     ...todo:: color all args, not just first
 251     """
 252     foreground = None
 253     background = None
 254     style = None
 255
 256     # remove color info from kwargs
 257     try:
 258         foreground = kwargs['foreground']
 259         del kwargs['foreground']
 260     except KeyError:
 261         pass
 262
 263     try:
 264         background = kwargs['background']
 265         del kwargs['background']
 266     except KeyError:
 267         pass
 268
 269     try:
 270         style = kwargs['style']
 271         del kwargs['style']
 272     except KeyError:
 273         pass
 274
 275     if _STDOUT_IS_TTY:
 276         text_c = colored(text, foreground, background, style)
 277     else:
 278         text_c = text
 279
 280     _builtin_print(text_c, *args, **kwargs)
 281
 282
 283 black   = partial(print, foreground=COLOR_BLACK)
 284 red     = partial(print, foreground=COLOR_RED)
 285 green   = partial(print, foreground=COLOR_GREEN)
 286 yellow  = partial(print, foreground=COLOR_YELLOW)
 287 blue    = partial(print, foreground=COLOR_BLUE)
 288 magenta = partial(print, foreground=COLOR_MAGENTA)
 289 cyan    = partial(print, foreground=COLOR_CYAN)
 290 white   = partial(print, foreground=COLOR_WHITE)
 291
 292 normal    = partial(print, style=STYLE_NORMAL)
 293 bold      = partial(print, style=STYLE_BOLD)
 294 underline = partial(print, style=STYLE_UNDERLINE)
 295 blink     = partial(print, style=STYLE_BLINK)
 296 reverse   = partial(print, style=STYLE_REVERSE)