src/pyutils/text_utils.py

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