2 # -*- coding: utf-8 -*-
4 # © Copyright 2021-2022, Scott Gasch
6 """Utilities for dealing with "text"."""
15 from collections import defaultdict
16 from dataclasses import dataclass
17 from typing import Dict, Generator, List, Literal, Optional, Tuple
20 from ansi import fg, reset
22 logger = logging.getLogger(__file__)
33 """Number of columns"""
36 def get_console_rows_columns() -> RowsColumns:
39 The number of rows/columns on the current console or None
40 if we can't tell or an error occurred.
42 from exec_utils import cmd
44 rows: Optional[str] = os.environ.get('LINES', None)
45 cols: Optional[str] = os.environ.get('COLUMNS', None)
46 if not rows or not cols:
74 if not rows or not cols:
75 raise Exception('Can\'t determine console size?!')
76 return RowsColumns(int(rows), int(cols))
79 class BarGraphText(enum.Enum):
80 """What kind of text to include at the end of the bar graph?"""
83 """None, leave it blank."""
97 text: BarGraphText = BarGraphText.PERCENTAGE,
98 fgcolor=fg("school bus yellow"),
103 """Draws a progress graph at the current cursor position.
106 current: how many have we done so far?
107 total: how many are there to do total?
108 text: how should we render the text at the end?
109 width: how many columns wide should be progress graph be?
110 fgcolor: what color should "done" part of the graph be?
111 left_end: the character at the left side of the graph
112 right_end: the character at the right side of the graph
113 redraw: if True, omit a line feed after the carriage return
114 so that subsequent calls to this method redraw the graph
117 ret = "\r" if redraw else "\n"
118 bar = bar_graph_string(
127 print(bar, end=ret, flush=True, file=sys.stderr)
130 def _make_bar_graph_text(text: BarGraphText, current: int, total: int, percentage: float):
131 if text == BarGraphText.NONE:
133 elif text == BarGraphText.PERCENTAGE:
134 return f'{percentage:.1f}'
135 elif text == BarGraphText.FRACTION:
136 return f'{current} / {total}'
137 raise ValueError(text)
140 def bar_graph_string(
144 text: BarGraphText = BarGraphText.PERCENTAGE,
146 fgcolor=fg("school bus yellow"),
151 """Returns a string containing a bar graph.
154 current: how many have we done so far?
155 total: how many are there to do total?
156 text: how should we render the text at the end?
157 width: how many columns wide should be progress graph be?
158 fgcolor: what color should "done" part of the graph be?
159 reset_seq: sequence to use to turn off color
160 left_end: the character at the left side of the graph
161 right_end: the character at the right side of the graph
163 >>> bar_graph(5, 10, fgcolor='', reset_seq='')
164 '[███████████████████████████████████ ] 50.0%'
169 percentage = float(current) / float(total)
172 if percentage < 0.0 or percentage > 1.0:
173 raise ValueError(percentage)
174 text = _make_bar_graph_text(text, current, total, percentage)
175 whole_width = math.floor(percentage * width)
176 if whole_width == width:
179 elif whole_width == 0 and percentage > 0.0:
182 remainder_width = (percentage * width) % 1
183 part_width = math.floor(remainder_width * 8)
184 part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
190 + " " * (width - whole_width - 1)
198 def sparkline(numbers: List[float]) -> Tuple[float, float, str]:
200 Makes a "sparkline" little inline histogram graph. Auto scales.
203 numbers: the population over which to create the sparkline
206 a three tuple containing:
208 * the minimum number in the population
209 * the maximum number in the population
210 * a string representation of the population in a concise format
212 >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7])
215 >>> sparkline([104, 99, 93, 96, 82, 77, 85, 73])
216 (73, 104, '█▇▆▆▃▂▄▁')
219 _bar = '▁▂▃▄▅▆▇█' # Unicode: 9601, 9602, 9603, 9604, 9605, 9606, 9607, 9608
222 min_num, max_num = min(numbers), max(numbers)
223 span = max_num - min_num
225 _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers
227 return min_num, max_num, sline
230 def distribute_strings(
237 Distributes strings into a line for justified text.
240 strings: a list of string tokens to distribute
241 width: the width of the line to create
242 padding: the padding character to place between string chunks
245 The distributed, justified string.
247 >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
250 ret = ' ' + ' '.join(strings) + ' '
251 assert len(string_utils.strip_ansi_sequences(ret)) < width
253 while len(string_utils.strip_ansi_sequences(ret)) < width:
254 spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
258 ret = before + padding + after
265 def _justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
267 Justifies a string chunk by chunk.
270 string: the string to be justified
271 width: how wide to make the output
272 padding: what padding character to use between chunks
277 >>> _justify_string_by_chunk("This is a test", 40)
279 >>> _justify_string_by_chunk("This is a test", 20)
283 assert len(string_utils.strip_ansi_sequences(string)) <= width
285 first, *rest, last = string.split()
287 len(string_utils.strip_ansi_sequences(first)) + len(string_utils.strip_ansi_sequences(last))
289 ret = first + distribute_strings(rest, width=w, padding=padding) + last
294 string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
296 """Justify a string to width with left, right, center of justified
300 string: the string to justify
301 width: the width to justify the string to
302 alignment: a single character indicating the desired alignment:
303 * 'c' = centered within the width
304 * 'j' = justified at width
305 * 'l' = left alignment
306 * 'r' = right alignment
307 padding: the padding character to use while justifying
309 >>> justify_string('This is another test', width=40, alignment='c')
310 ' This is another test '
311 >>> justify_string('This is another test', width=40, alignment='l')
312 'This is another test '
313 >>> justify_string('This is another test', width=40, alignment='r')
314 ' This is another test'
315 >>> justify_string('This is another test', width=40, alignment='j')
316 'This is another test'
318 alignment = alignment[0]
320 while len(string_utils.strip_ansi_sequences(string)) < width:
323 elif alignment == "r":
324 string = padding + string
325 elif alignment == "j":
326 return _justify_string_by_chunk(string, width=width, padding=padding)
327 elif alignment == "c":
328 if len(string) % 2 == 0:
331 string = padding + string
337 def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by: int = 0) -> str:
338 """Justifies text with left, right, centered or justified alignment
339 and optionally with initial indentation.
342 text: the text to be justified
343 width: the width at which to justify text
344 alignment: a single character indicating the desired alignment:
345 * 'c' = centered within the width
346 * 'j' = justified at width
347 * 'l' = left alignment
348 * 'r' = right alignment
349 indent_by: if non-zero, adds n prefix spaces to indent the text.
354 >>> justify_text('This is a test of the emergency broadcast system. This is only a test.',
355 ... width=40, alignment='j') #doctest: +NORMALIZE_WHITESPACE
356 'This is a test of the emergency\\nbroadcast system. This is only a test.'
362 indent += ' ' * indent_by
365 for word in text.split():
367 len(string_utils.strip_ansi_sequences(line))
368 + len(string_utils.strip_ansi_sequences(word))
371 line = justify_string(line, width=width, alignment=alignment)
372 retval = retval + '\n' + line
374 line = line + ' ' + word
375 if len(string_utils.strip_ansi_sequences(line)) > 0:
377 retval += "\n" + justify_string(line[1:], width=width, alignment=alignment)
379 retval += "\n" + line[1:]
383 def generate_padded_columns(text: List[str]) -> Generator:
384 """Given a list of strings, break them into columns using :meth:`split`
385 and then compute the maximum width of each column. Finally,
386 distribute the columular chunks into the output padding each to
390 text: a list of strings to chunk into padded columns
393 padded columns based on text.split()
395 >>> for x in generate_padded_columns(
396 ... [ 'reading writing arithmetic',
397 ... 'mathematics psychology physics',
398 ... 'communications sociology anthropology' ]):
400 reading writing arithmetic
401 mathematics psychology physics
402 communications sociology anthropology
404 max_width: Dict[int, int] = defaultdict(int)
406 for pos, word in enumerate(line.split()):
407 max_width[pos] = max(max_width[pos], len(string_utils.strip_ansi_sequences(word)))
411 for pos, word in enumerate(line.split()):
412 width = max_width[pos]
413 word = justify_string(word, width=width, alignment='l')
418 def wrap_string(text: str, n: int) -> str:
421 text: the string to be wrapped
422 n: the width after which to wrap text
425 The wrapped form of text
427 chunks = text.split()
431 if width + len(string_utils.strip_ansi_sequences(chunk)) > n:
435 width += len(string_utils.strip_ansi_sequences(chunk)) + 1
439 class Indenter(contextlib.AbstractContextManager):
441 Context manager that indents stuff (even recursively). e.g.::
443 with Indenter(pad_count = 8) as i:
460 pad_prefix: Optional[str] = None,
464 """Construct an Indenter.
467 pad_prefix: an optional prefix to prepend to each line
468 pad_char: the character used to indent
469 pad_count: the number of pad_chars to use to indent
472 if pad_prefix is not None:
473 self.pad_prefix = pad_prefix
476 self.padding = pad_char * pad_count
482 def __exit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
488 def print(self, *arg, **kwargs):
489 text = string_utils.sprintf(*arg, **kwargs)
490 print(self.pad_prefix + self.padding * self.level + text, end='')
496 width: Optional[int] = None,
497 align: Optional[str] = None,
498 style: Optional[str] = 'solid',
499 color: Optional[str] = None,
502 Creates a nice header line with a title.
506 width: how wide to make the header
507 align: "left" or "right"
508 style: "ascii", "solid" or "dashed"
511 The header as a string.
513 >>> header('title', width=60, style='ascii')
514 '----[ title ]-----------------------------------------------'
518 width = get_console_rows_columns().columns
526 text_len = len(string_utils.strip_ansi_sequences(title))
529 right = width - (left + text_len + 4)
530 elif align == 'right':
532 left = width - (right + text_len + 4)
534 left = int((width - (text_len + 4)) / 2)
536 while left + text_len + 4 + right < width:
543 elif style == 'dashed':
557 return line_char * left + begin + col + ' ' + title + ' ' + reset_seq + end + line_char * right
561 title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
564 Make a nice unicode box (optionally with color) around some text.
567 title: the title of the box
568 text: the text in the box
569 width: the box's width
570 color: the box's color
575 >>> print(box('title', 'this is some text', width=20).strip())
585 text = justify_text(text, width=width - 4, alignment='l')
586 return preformatted_box(title, text, width=width, color=color)
589 def preformatted_box(
590 title: Optional[str] = None, text: Optional[str] = None, *, width=80, color: str = ''
592 """Creates a nice box with rounded corners and returns it as a string.
595 title: the title of the box
596 text: the text inside the box
597 width: the width of the box
598 color: the box's color
603 >>> print(preformatted_box('title', 'this\\nis\\nsome\\ntext', width=20).strip())
620 ret += color + '╭' + '─' * w + '╮' + rset + '\n'
621 if title is not None:
626 + justify_string(title, width=w, alignment='c')
632 ret += color + '│' + ' ' * w + '│' + rset + '\n'
634 for line in text.split('\n'):
635 tw = len(string_utils.strip_ansi_sequences(line))
637 ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
638 ret += color + '╰' + '─' * w + '╯' + rset + '\n'
643 title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
645 """Draws a box with nice rounded corners.
647 >>> print_box('Title', 'This is text', width=30)
648 ╭────────────────────────────╮
652 ╰────────────────────────────╯
654 >>> print_box(None, 'OK', width=6)
659 print(preformatted_box(title, text, width=width, color=color), end='')
662 if __name__ == '__main__':