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