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