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