Commit | Line | Data |
---|---|---|
f365f614 CH |
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> | |
c16165f4 | 20 | |
856fe417 | 21 | """ |
80648d47 CH |
22 | Functions for improving textual output. |
23 | ||
24 | Copyright: 2015 Intra2net AG | |
25 | ||
80648d47 | 26 | This module has two parts. Part 1 includes: |
7628bc48 CH |
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 | |
80648d47 CH |
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) | |
c16165f4 CH |
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 | ||
856fe417 | 42 | |
80648d47 CH |
43 | .. seealso:: http://stackoverflow.com/questions/287871/print-in-terminal-with-colors-using-python |
44 | .. seealso:: https://en.wikipedia.org/wiki/ANSI_escape_code | |
7f66ff3e | 45 | .. seealso:: :py:mod:`textwrap` |
c16165f4 CH |
46 | """ |
47 | ||
7628bc48 | 48 | from builtins import print as _builtin_print |
85fe93ec | 49 | |
c16165f4 | 50 | from functools import partial |
d8502289 | 51 | from itertools import islice |
c16165f4 | 52 | |
1242b1cf | 53 | from .type_helpers import isstr |
170db03f | 54 | from sys import stdout |
d8502289 CH |
55 | |
56 | ||
eb048121 | 57 | def head_and_tail(iterable, n_head=20, n_tail=20, n_elems=None, |
170db03f | 58 | skip_elem="...(skipping {n_skipped} elements)..."): |
80648d47 CH |
59 | """ |
60 | Convenient way to shorten a possibly very long iterable before printing. | |
d8502289 CH |
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 | |
80648d47 | 67 | iterables this will be faster if random element access is provided via [] |
d8502289 CH |
68 | |
69 | :param iterable: an iterable | |
80648d47 | 70 | :type iterable: anything that can be iterated over |
d8502289 CH |
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 | |
80648d47 CH |
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). | |
d8502289 | 82 | |
7f66ff3e CH |
83 | .. seealso:: :py:func:`slice`, :py:func:`itertools.islice`, :py:func:`textwrap.shorten` |
84 | ||
d8502289 | 85 | """ |
d8502289 CH |
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: | |
170db03f | 100 | if isstr(skip_elem): |
d8502289 CH |
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: | |
d6f2f64e | 105 | return |
d8502289 CH |
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 | ||
856fe417 | 121 | def size_str(byte_number, is_diff=False): |
80648d47 CH |
122 | """ |
123 | Create a human-readable text representation of a file size. | |
856fe417 | 124 | |
80648d47 CH |
125 | Rounds and shortens size to something easily human-readable like '1.5 GB'. |
126 | ||
7628bc48 | 127 | :param float byte_number: Number of bytes to express as text |
80648d47 | 128 | :param bool is_diff: Set to True to include a '+' or '-' in output; |
856fe417 | 129 | default: False |
80648d47 CH |
130 | :returns: textual representation of data |
131 | :rtype: str | |
856fe417 | 132 | """ |
856fe417 CH |
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 = '' | |
856fe417 CH |
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 | |
7628bc48 | 159 | return f'{sign_str}{int(round(curr_num*factor)):,d} {units[-1]}B' |
856fe417 CH |
160 | |
161 | ||
d8502289 CH |
162 | ############################################################################### |
163 | # TEXT FORMATTING/COLORING | |
164 | ############################################################################### | |
c16165f4 CH |
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 | ||
7628bc48 | 182 | _COLOR_TO_CODE = dict(zip([COLOR_BLACK, COLOR_RED, COLOR_GREEN, COLOR_YELLOW, |
c16165f4 CH |
183 | COLOR_BLUE, COLOR_MAGENTA, COLOR_CYAN, COLOR_WHITE], |
184 | range(8))) | |
185 | ||
186 | _ANSI_ESCAPE_SURROUND = '\x1b[{}m{}\x1b[0m' | |
187 | ||
49439cda CH |
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 | |
170db03f | 194 | |
c16165f4 CH |
195 | |
196 | def colored(text, foreground=None, background=None, style=None): | |
7628bc48 | 197 | """ return text with given foreground/background ANSI color escape sequence |
c16165f4 CH |
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 | |
7628bc48 | 203 | :param str background: one of the color constants to use for background |
c16165f4 CH |
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 | ||
170db03f CH |
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 | |
d5d2c1d7 | 232 | color-args through that function. Runs built-in :py:func:`print` |
170db03f CH |
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) | |
c16165f4 CH |
266 | |
267 | ||
170db03f CH |
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) | |
c16165f4 | 273 | magenta = partial(print, foreground=COLOR_MAGENTA) |
170db03f CH |
274 | cyan = partial(print, foreground=COLOR_CYAN) |
275 | white = partial(print, foreground=COLOR_WHITE) | |
c16165f4 CH |
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) |