text_utils.py

   1 #!/usr/bin/env python3
   2 # -*- coding: utf-8 -*-
   3
   4 # © Copyright 2021-2022, Scott Gasch
   5
   6 """Utilities for dealing with "text"."""
   7
   8 import contextlib
   9 import logging
  10 import math
  11 import os
  12 import re
  13 import sys
  14 from collections import defaultdict
  15 from dataclasses import dataclass
  16 from typing import Dict, Generator, List, Literal, Optional, Tuple
  17
  18 import string_utils
  19 from ansi import fg, reset
  20
  21 logger = logging.getLogger(__file__)
  22
  23
  24 @dataclass
  25 class RowsColumns:
  26     """Row + Column"""
  27
  28     rows: int = 0
  29     """Numer of rows"""
  30
  31     columns: int = 0
  32     """Number of columns"""
  33
  34
  35 def get_console_rows_columns() -> RowsColumns:
  36     """
  37     Returns:
  38         The number of rows/columns on the current console or None
  39         if we can't tell or an error occurred.
  40     """
  41     from exec_utils import cmd
  42
  43     rows: Optional[str] = os.environ.get('LINES', None)
  44     cols: Optional[str] = os.environ.get('COLUMNS', None)
  45     if not rows or not cols:
  46         try:
  47             rows, cols = cmd(
  48                 "stty size",
  49                 timeout_seconds=1.0,
  50             ).split()
  51         except Exception:
  52             rows = None
  53             cols = None
  54
  55     if rows is None:
  56         try:
  57             rows = cmd(
  58                 "tput rows",
  59                 timeout_seconds=1.0,
  60             )
  61         except Exception:
  62             rows = None
  63
  64     if cols is None:
  65         try:
  66             cols = cmd(
  67                 "tput cols",
  68                 timeout_seconds=1.0,
  69             )
  70         except Exception:
  71             cols = None
  72
  73     if not rows or not cols:
  74         raise Exception('Can\'t determine console size?!')
  75     return RowsColumns(int(rows), int(cols))
  76
  77
  78 def progress_graph(
  79     current: int,
  80     total: int,
  81     *,
  82     width=70,
  83     fgcolor=fg("school bus yellow"),
  84     left_end="[",
  85     right_end="]",
  86     redraw=True,
  87 ) -> None:
  88     """Draws a progress graph at the current cursor position.
  89
  90     Args:
  91         current: how many have we done so far?
  92         total: how many are there to do total?
  93         width: how many columns wide should be progress graph be?
  94         fgcolor: what color should "done" part of the graph be?
  95         left_end: the character at the left side of the graph
  96         right_end: the character at the right side of the graph
  97         redraw: if True, omit a line feed after the carriage return
  98             so that subsequent calls to this method redraw the graph
  99             iteratively.
 100     """
 101     percent = current / total
 102     ret = "\r" if redraw else "\n"
 103     bar = bar_graph(
 104         percent,
 105         include_text=True,
 106         width=width,
 107         fgcolor=fgcolor,
 108         left_end=left_end,
 109         right_end=right_end,
 110     )
 111     print(bar, end=ret, flush=True, file=sys.stderr)
 112
 113
 114 def bar_graph(
 115     percentage: float,
 116     *,
 117     include_text=True,
 118     width=70,
 119     fgcolor=fg("school bus yellow"),
 120     reset_seq=reset(),
 121     left_end="[",
 122     right_end="]",
 123 ) -> str:
 124     """Returns a string containing a bar graph.
 125
 126     Args:
 127         percentage: percentage complete (0..100)
 128         include_text: should we include the percentage text at the end?
 129         width: how many columns wide should be progress graph be?
 130         fgcolor: what color should "done" part of the graph be?
 131         reset_seq: sequence to use to turn off color
 132         left_end: the character at the left side of the graph
 133         right_end: the character at the right side of the graph
 134
 135     >>> bar_graph(0.5, fgcolor='', reset_seq='')
 136     '[███████████████████████████████████                                   ] 50.0%'
 137
 138     """
 139
 140     if percentage < 0.0 or percentage > 1.0:
 141         raise ValueError(percentage)
 142     if include_text:
 143         text = f"{percentage*100.0:2.1f}%"
 144     else:
 145         text = ""
 146     whole_width = math.floor(percentage * width)
 147     if whole_width == width:
 148         whole_width -= 1
 149         part_char = "▉"
 150     elif whole_width == 0 and percentage > 0.0:
 151         part_char = "▏"
 152     else:
 153         remainder_width = (percentage * width) % 1
 154         part_width = math.floor(remainder_width * 8)
 155         part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
 156     return (
 157         left_end
 158         + fgcolor
 159         + "█" * whole_width
 160         + part_char
 161         + " " * (width - whole_width - 1)
 162         + reset_seq
 163         + right_end
 164         + " "
 165         + text
 166     )
 167
 168
 169 def sparkline(numbers: List[float]) -> Tuple[float, float, str]:
 170     """
 171     Makes a "sparkline" little inline histogram graph.  Auto scales.
 172
 173     Args:
 174         numbers: the population over which to create the sparkline
 175
 176     Returns:
 177         a three tuple containing:
 178
 179         * the minimum number in the population
 180         * the maximum number in the population
 181         * a string representation of the population in a concise format
 182
 183     >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7])
 184     (1, 10, '▁▁▂▄█▂▄▆')
 185
 186     >>> sparkline([104, 99, 93, 96, 82, 77, 85, 73])
 187     (73, 104, '█▇▆▆▃▂▄▁')
 188
 189     """
 190     _bar = '▁▂▃▄▅▆▇█'  # Unicode: 9601, 9602, 9603, 9604, 9605, 9606, 9607, 9608
 191
 192     barcount = len(_bar)
 193     min_num, max_num = min(numbers), max(numbers)
 194     span = max_num - min_num
 195     sline = ''.join(
 196         _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers
 197     )
 198     return min_num, max_num, sline
 199
 200
 201 def distribute_strings(
 202     strings: List[str],
 203     *,
 204     width: int = 80,
 205     padding: str = " ",
 206 ) -> str:
 207     """
 208     Distributes strings into a line for justified text.
 209
 210     Args:
 211         strings: a list of string tokens to distribute
 212         width: the width of the line to create
 213         padding: the padding character to place between string chunks
 214
 215     Returns:
 216         The distributed, justified string.
 217
 218     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
 219     '      this      is      a      test     '
 220     """
 221     ret = ' ' + ' '.join(strings) + ' '
 222     assert len(string_utils.strip_ansi_sequences(ret)) < width
 223     x = 0
 224     while len(string_utils.strip_ansi_sequences(ret)) < width:
 225         spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
 226         where = spaces[x]
 227         before = ret[:where]
 228         after = ret[where:]
 229         ret = before + padding + after
 230         x += 1
 231         if x >= len(spaces):
 232             x = 0
 233     return ret
 234
 235
 236 def _justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
 237     """
 238     Justifies a string chunk by chunk.
 239
 240     Args:
 241         string: the string to be justified
 242         width: how wide to make the output
 243         padding: what padding character to use between chunks
 244
 245     Returns:
 246         the justified string
 247
 248     >>> _justify_string_by_chunk("This is a test", 40)
 249     'This          is          a         test'
 250     >>> _justify_string_by_chunk("This is a test", 20)
 251     'This   is   a   test'
 252
 253     """
 254     assert len(string_utils.strip_ansi_sequences(string)) <= width
 255     padding = padding[0]
 256     first, *rest, last = string.split()
 257     w = width - (
 258         len(string_utils.strip_ansi_sequences(first)) + len(string_utils.strip_ansi_sequences(last))
 259     )
 260     ret = first + distribute_strings(rest, width=w, padding=padding) + last
 261     return ret
 262
 263
 264 def justify_string(
 265     string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
 266 ) -> str:
 267     """Justify a string to width with left, right, center of justified
 268     alignment.
 269
 270     Args:
 271         string: the string to justify
 272         width: the width to justify the string to
 273         alignment: a single character indicating the desired alignment:
 274             * 'c' = centered within the width
 275             * 'j' = justified at width
 276             * 'l' = left alignment
 277             * 'r' = right alignment
 278         padding: the padding character to use while justifying
 279
 280     >>> justify_string('This is another test', width=40, alignment='c')
 281     '          This is another test          '
 282     >>> justify_string('This is another test', width=40, alignment='l')
 283     'This is another test                    '
 284     >>> justify_string('This is another test', width=40, alignment='r')
 285     '                    This is another test'
 286     >>> justify_string('This is another test', width=40, alignment='j')
 287     'This        is        another       test'
 288     """
 289     alignment = alignment[0]
 290     padding = padding[0]
 291     while len(string_utils.strip_ansi_sequences(string)) < width:
 292         if alignment == "l":
 293             string += padding
 294         elif alignment == "r":
 295             string = padding + string
 296         elif alignment == "j":
 297             return _justify_string_by_chunk(string, width=width, padding=padding)
 298         elif alignment == "c":
 299             if len(string) % 2 == 0:
 300                 string += padding
 301             else:
 302                 string = padding + string
 303         else:
 304             raise ValueError
 305     return string
 306
 307
 308 def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by: int = 0) -> str:
 309     """Justifies text with left, right, centered or justified alignment
 310     and optionally with initial indentation.
 311
 312     Args:
 313         text: the text to be justified
 314         width: the width at which to justify text
 315         alignment: a single character indicating the desired alignment:
 316             * 'c' = centered within the width
 317             * 'j' = justified at width
 318             * 'l' = left alignment
 319             * 'r' = right alignment
 320         indent_by: if non-zero, adds n prefix spaces to indent the text.
 321
 322     Returns:
 323         The justified text.
 324
 325     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
 326     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
 327     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 328
 329     """
 330     retval = ''
 331     indent = ''
 332     if indent_by > 0:
 333         indent += ' ' * indent_by
 334     line = indent
 335
 336     for word in text.split():
 337         if (
 338             len(string_utils.strip_ansi_sequences(line))
 339             + len(string_utils.strip_ansi_sequences(word))
 340         ) > width:
 341             line = line[1:]
 342             line = justify_string(line, width=width, alignment=alignment)
 343             retval = retval + '\n' + line
 344             line = indent
 345         line = line + ' ' + word
 346     if len(string_utils.strip_ansi_sequences(line)) > 0:
 347         if alignment != 'j':
 348             retval += "\n" + justify_string(line[1:], width=width, alignment=alignment)
 349         else:
 350             retval += "\n" + line[1:]
 351     return retval[1:]
 352
 353
 354 def generate_padded_columns(text: List[str]) -> Generator:
 355     """Given a list of strings, break them into columns using :meth:`split`
 356     and then compute the maximum width of each column.  Finally,
 357     distribute the columular chunks into the output padding each to
 358     the proper width.
 359
 360     Args:
 361         text: a list of strings to chunk into padded columns
 362
 363     Returns:
 364         padded columns based on text.split()
 365
 366     >>> for x in generate_padded_columns(
 367     ...     [ 'reading writing arithmetic',
 368     ...       'mathematics psychology physics',
 369     ...       'communications sociology anthropology' ]):
 370     ...     print(x.strip())
 371     reading        writing    arithmetic
 372     mathematics    psychology physics
 373     communications sociology  anthropology
 374     """
 375     max_width: Dict[int, int] = defaultdict(int)
 376     for line in text:
 377         for pos, word in enumerate(line.split()):
 378             max_width[pos] = max(max_width[pos], len(string_utils.strip_ansi_sequences(word)))
 379
 380     for line in text:
 381         out = ""
 382         for pos, word in enumerate(line.split()):
 383             width = max_width[pos]
 384             word = justify_string(word, width=width, alignment='l')
 385             out += f'{word} '
 386         yield out
 387
 388
 389 def wrap_string(text: str, n: int) -> str:
 390     """
 391     Args:
 392         text: the string to be wrapped
 393         n: the width after which to wrap text
 394
 395     Returns:
 396         The wrapped form of text
 397     """
 398     chunks = text.split()
 399     out = ''
 400     width = 0
 401     for chunk in chunks:
 402         if width + len(string_utils.strip_ansi_sequences(chunk)) > n:
 403             out += '\n'
 404             width = 0
 405         out += chunk + ' '
 406         width += len(string_utils.strip_ansi_sequences(chunk)) + 1
 407     return out
 408
 409
 410 class Indenter(contextlib.AbstractContextManager):
 411     """
 412     Context manager that indents stuff (even recursively).  e.g.::
 413
 414         with Indenter(pad_count = 8) as i:
 415             i.print('test')
 416             with i:
 417                 i.print('-ing')
 418                 with i:
 419                     i.print('1, 2, 3')
 420
 421     Yields::
 422
 423         test
 424                 -ing
 425                         1, 2, 3
 426     """
 427
 428     def __init__(
 429         self,
 430         *,
 431         pad_prefix: Optional[str] = None,
 432         pad_char: str = ' ',
 433         pad_count: int = 4,
 434     ):
 435         """Construct an Indenter.
 436
 437         Args:
 438             pad_prefix: an optional prefix to prepend to each line
 439             pad_char: the character used to indent
 440             pad_count: the number of pad_chars to use to indent
 441         """
 442         self.level = -1
 443         if pad_prefix is not None:
 444             self.pad_prefix = pad_prefix
 445         else:
 446             self.pad_prefix = ''
 447         self.padding = pad_char * pad_count
 448
 449     def __enter__(self):
 450         self.level += 1
 451         return self
 452
 453     def __exit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
 454         self.level -= 1
 455         if self.level < -1:
 456             self.level = -1
 457         return False
 458
 459     def print(self, *arg, **kwargs):
 460         text = string_utils.sprintf(*arg, **kwargs)
 461         print(self.pad_prefix + self.padding * self.level + text, end='')
 462
 463
 464 def header(
 465     title: str,
 466     *,
 467     width: Optional[int] = None,
 468     align: Optional[str] = None,
 469     style: Optional[str] = 'solid',
 470     color: Optional[str] = None,
 471 ):
 472     """
 473     Creates a nice header line with a title.
 474
 475     Args:
 476         title: the title
 477         width: how wide to make the header
 478         align: "left" or "right"
 479         style: "ascii", "solid" or "dashed"
 480
 481     Returns:
 482         The header as a string.
 483
 484     >>> header('title', width=60, style='ascii')
 485     '----[ title ]-----------------------------------------------'
 486     """
 487     if not width:
 488         try:
 489             width = get_console_rows_columns().columns
 490         except Exception:
 491             width = 80
 492     if not align:
 493         align = 'left'
 494     if not style:
 495         style = 'ascii'
 496
 497     text_len = len(string_utils.strip_ansi_sequences(title))
 498     if align == 'left':
 499         left = 4
 500         right = width - (left + text_len + 4)
 501     elif align == 'right':
 502         right = 4
 503         left = width - (right + text_len + 4)
 504     else:
 505         left = int((width - (text_len + 4)) / 2)
 506         right = left
 507         while left + text_len + 4 + right < width:
 508             right += 1
 509
 510     if style == 'solid':
 511         line_char = '━'
 512         begin = ''
 513         end = ''
 514     elif style == 'dashed':
 515         line_char = '┅'
 516         begin = ''
 517         end = ''
 518     else:
 519         line_char = '-'
 520         begin = '['
 521         end = ']'
 522     if color:
 523         col = color
 524         reset_seq = reset()
 525     else:
 526         col = ''
 527         reset_seq = ''
 528     return line_char * left + begin + col + ' ' + title + ' ' + reset_seq + end + line_char * right
 529
 530
 531 def box(
 532     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 533 ) -> str:
 534     """
 535     Make a nice unicode box (optionally with color) around some text.
 536
 537     Args:
 538         title: the title of the box
 539         text: the text in the box
 540         width: the box's width
 541         color: the box's color
 542
 543     Returns:
 544         the box as a string
 545
 546     >>> print(box('title', 'this is some text', width=20).strip())
 547     ╭──────────────────╮
 548     │       title      │
 549     │                  │
 550     │ this is some     │
 551     │ text             │
 552     ╰──────────────────╯
 553     """
 554     assert width > 4
 555     if text is not None:
 556         text = justify_text(text, width=width - 4, alignment='l')
 557     return preformatted_box(title, text, width=width, color=color)
 558
 559
 560 def preformatted_box(
 561     title: Optional[str] = None, text: Optional[str] = None, *, width=80, color: str = ''
 562 ) -> str:
 563     """Creates a nice box with rounded corners and returns it as a string.
 564
 565     Args:
 566         title: the title of the box
 567         text: the text inside the box
 568         width: the width of the box
 569         color: the box's color
 570
 571     Returns:
 572         the box as a string
 573
 574     >>> print(preformatted_box('title', 'this\\nis\\nsome\\ntext', width=20).strip())
 575     ╭──────────────────╮
 576     │       title      │
 577     │                  │
 578     │ this             │
 579     │ is               │
 580     │ some             │
 581     │ text             │
 582     ╰──────────────────╯
 583     """
 584     assert width > 4
 585     ret = ''
 586     if color == '':
 587         rset = ''
 588     else:
 589         rset = reset()
 590     w = width - 2
 591     ret += color + '╭' + '─' * w + '╮' + rset + '\n'
 592     if title is not None:
 593         ret += (
 594             color
 595             + '│'
 596             + rset
 597             + justify_string(title, width=w, alignment='c')
 598             + color
 599             + '│'
 600             + rset
 601             + '\n'
 602         )
 603         ret += color + '│' + ' ' * w + '│' + rset + '\n'
 604     if text is not None:
 605         for line in text.split('\n'):
 606             tw = len(string_utils.strip_ansi_sequences(line))
 607             assert tw <= w
 608             ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
 609     ret += color + '╰' + '─' * w + '╯' + rset + '\n'
 610     return ret
 611
 612
 613 def print_box(
 614     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 615 ) -> None:
 616     """Draws a box with nice rounded corners.
 617
 618     >>> print_box('Title', 'This is text', width=30)
 619     ╭────────────────────────────╮
 620     │            Title           │
 621     │                            │
 622     │ This is text               │
 623     ╰────────────────────────────╯
 624
 625     >>> print_box(None, 'OK', width=6)
 626     ╭────╮
 627     │ OK │
 628     ╰────╯
 629     """
 630     print(preformatted_box(title, text, width=width, color=color), end='')
 631
 632
 633 if __name__ == '__main__':
 634     import doctest
 635
 636     doctest.testmod()