src/pyutils/text_utils.py

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