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