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