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