src/pyutils/text_utils.py

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