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