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