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 logging
  10 import math
  11 import os
  12 import re
  13 import sys
  14 from collections import defaultdict
  15 from dataclasses import dataclass
  16 from typing import Dict, Generator, List, Literal, Optional, Tuple
  17
  18 import string_utils
  19 from ansi import fg, reset
  20
  21 logger = logging.getLogger(__file__)
  22
  23
  24 @dataclass
  25 class RowsColumns:
  26     """Row + Column"""
  27
  28     rows: int = 0
  29     columns: int = 0
  30
  31
  32 def get_console_rows_columns() -> RowsColumns:
  33     """Returns the number of rows/columns on the current console."""
  34
  35     from exec_utils import cmd
  36
  37     rows: Optional[str] = os.environ.get('LINES', None)
  38     cols: Optional[str] = os.environ.get('COLUMNS', None)
  39     if not rows or not cols:
  40         try:
  41             rows, cols = cmd(
  42                 "stty size",
  43                 timeout_seconds=1.0,
  44             ).split()
  45         except Exception:
  46             rows = None
  47             cols = None
  48
  49     if rows is None:
  50         try:
  51             rows = cmd(
  52                 "tput rows",
  53                 timeout_seconds=1.0,
  54             )
  55         except Exception:
  56             rows = None
  57
  58     if cols is None:
  59         try:
  60             cols = cmd(
  61                 "tput cols",
  62                 timeout_seconds=1.0,
  63             )
  64         except Exception:
  65             cols = None
  66
  67     if not rows or not cols:
  68         raise Exception('Can\'t determine console size?!')
  69     return RowsColumns(int(rows), int(cols))
  70
  71
  72 def progress_graph(
  73     current: int,
  74     total: int,
  75     *,
  76     width=70,
  77     fgcolor=fg("school bus yellow"),
  78     left_end="[",
  79     right_end="]",
  80     redraw=True,
  81 ) -> None:
  82     """Draws a progress graph."""
  83
  84     percent = current / total
  85     ret = "\r" if redraw else "\n"
  86     bar = bar_graph(
  87         percent,
  88         include_text=True,
  89         width=width,
  90         fgcolor=fgcolor,
  91         left_end=left_end,
  92         right_end=right_end,
  93     )
  94     print(bar, end=ret, flush=True, file=sys.stderr)
  95
  96
  97 def bar_graph(
  98     percentage: float,
  99     *,
 100     include_text=True,
 101     width=70,
 102     fgcolor=fg("school bus yellow"),
 103     reset_seq=reset(),
 104     left_end="[",
 105     right_end="]",
 106 ) -> str:
 107     """Returns a string containing a bar graph.
 108
 109     >>> bar_graph(0.5, fgcolor='', reset_seq='')
 110     '[███████████████████████████████████                                   ] 50.0%'
 111
 112     """
 113
 114     if percentage < 0.0 or percentage > 1.0:
 115         raise ValueError(percentage)
 116     if include_text:
 117         text = f"{percentage*100.0:2.1f}%"
 118     else:
 119         text = ""
 120     whole_width = math.floor(percentage * width)
 121     if whole_width == width:
 122         whole_width -= 1
 123         part_char = "▉"
 124     elif whole_width == 0 and percentage > 0.0:
 125         part_char = "▏"
 126     else:
 127         remainder_width = (percentage * width) % 1
 128         part_width = math.floor(remainder_width * 8)
 129         part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
 130     return (
 131         left_end
 132         + fgcolor
 133         + "█" * whole_width
 134         + part_char
 135         + " " * (width - whole_width - 1)
 136         + reset_seq
 137         + right_end
 138         + " "
 139         + text
 140     )
 141
 142
 143 def sparkline(numbers: List[float]) -> Tuple[float, float, str]:
 144     """
 145     Makes a "sparkline" little inline histogram graph.  Auto scales.
 146
 147     >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7])
 148     (1, 10, '▁▁▂▄█▂▄▆')
 149
 150     >>> sparkline([104, 99, 93, 96, 82, 77, 85, 73])
 151     (73, 104, '█▇▆▆▃▂▄▁')
 152
 153     """
 154     _bar = '▁▂▃▄▅▆▇█'  # Unicode: 9601, 9602, 9603, 9604, 9605, 9606, 9607, 9608
 155
 156     barcount = len(_bar)
 157     min_num, max_num = min(numbers), max(numbers)
 158     span = max_num - min_num
 159     sline = ''.join(
 160         _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers
 161     )
 162     return min_num, max_num, sline
 163
 164
 165 def distribute_strings(
 166     strings: List[str],
 167     *,
 168     width: int = 80,
 169     padding: str = " ",
 170 ) -> str:
 171     """
 172     Distributes strings into a line for justified text.
 173
 174     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
 175     '      this      is      a      test     '
 176
 177     """
 178     ret = ' ' + ' '.join(strings) + ' '
 179     assert len(string_utils.strip_ansi_sequences(ret)) < width
 180     x = 0
 181     while len(string_utils.strip_ansi_sequences(ret)) < width:
 182         spaces = [m.start() for m in re.finditer(r' ([^ ]|$)', ret)]
 183         where = spaces[x]
 184         before = ret[:where]
 185         after = ret[where:]
 186         ret = before + padding + after
 187         x += 1
 188         if x >= len(spaces):
 189             x = 0
 190     return ret
 191
 192
 193 def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
 194     """
 195     Justifies a string.
 196
 197     >>> justify_string_by_chunk("This is a test", 40)
 198     'This          is          a         test'
 199     >>> justify_string_by_chunk("This is a test", 20)
 200     'This   is   a   test'
 201
 202     """
 203     assert len(string_utils.strip_ansi_sequences(string)) <= width
 204     padding = padding[0]
 205     first, *rest, last = string.split()
 206     w = width - (
 207         len(string_utils.strip_ansi_sequences(first)) + len(string_utils.strip_ansi_sequences(last))
 208     )
 209     ret = first + distribute_strings(rest, width=w, padding=padding) + last
 210     return ret
 211
 212
 213 def justify_string(
 214     string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
 215 ) -> str:
 216     """Justify a string.
 217
 218     >>> justify_string('This is another test', width=40, alignment='c')
 219     '          This is another test          '
 220     >>> justify_string('This is another test', width=40, alignment='l')
 221     'This is another test                    '
 222     >>> justify_string('This is another test', width=40, alignment='r')
 223     '                    This is another test'
 224     >>> justify_string('This is another test', width=40, alignment='j')
 225     'This        is        another       test'
 226
 227     """
 228     alignment = alignment[0]
 229     padding = padding[0]
 230     while len(string_utils.strip_ansi_sequences(string)) < width:
 231         if alignment == "l":
 232             string += padding
 233         elif alignment == "r":
 234             string = padding + string
 235         elif alignment == "j":
 236             return justify_string_by_chunk(string, width=width, padding=padding)
 237         elif alignment == "c":
 238             if len(string) % 2 == 0:
 239                 string += padding
 240             else:
 241                 string = padding + string
 242         else:
 243             raise ValueError
 244     return string
 245
 246
 247 def justify_text(text: str, *, width: int = 80, alignment: str = "c", indent_by: int = 0) -> str:
 248     """
 249     Justifies text optionally with initial indentation.
 250
 251     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
 252     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
 253     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 254
 255     """
 256     retval = ''
 257     indent = ''
 258     if indent_by > 0:
 259         indent += ' ' * indent_by
 260     line = indent
 261
 262     for word in text.split():
 263         if (
 264             len(string_utils.strip_ansi_sequences(line))
 265             + len(string_utils.strip_ansi_sequences(word))
 266         ) > width:
 267             line = line[1:]
 268             line = justify_string(line, width=width, alignment=alignment)
 269             retval = retval + '\n' + line
 270             line = indent
 271         line = line + ' ' + word
 272     if len(string_utils.strip_ansi_sequences(line)) > 0:
 273         if alignment != 'j':
 274             retval += "\n" + justify_string(line[1:], width=width, alignment=alignment)
 275         else:
 276             retval += "\n" + line[1:]
 277     return retval[1:]
 278
 279
 280 def generate_padded_columns(text: List[str]) -> Generator:
 281     max_width: Dict[int, int] = defaultdict(int)
 282     for line in text:
 283         for pos, word in enumerate(line.split()):
 284             max_width[pos] = max(max_width[pos], len(string_utils.strip_ansi_sequences(word)))
 285
 286     for line in text:
 287         out = ""
 288         for pos, word in enumerate(line.split()):
 289             width = max_width[pos]
 290             word = justify_string(word, width=width, alignment='l')
 291             out += f'{word} '
 292         yield out
 293
 294
 295 def wrap_string(text: str, n: int) -> str:
 296     chunks = text.split()
 297     out = ''
 298     width = 0
 299     for chunk in chunks:
 300         if width + len(string_utils.strip_ansi_sequences(chunk)) > n:
 301             out += '\n'
 302             width = 0
 303         out += chunk + ' '
 304         width += len(string_utils.strip_ansi_sequences(chunk)) + 1
 305     return out
 306
 307
 308 class Indenter(contextlib.AbstractContextManager):
 309     """
 310     Context manager that indents stuff (even recursively).  e.g.::
 311
 312         with Indenter(pad_count = 8) as i:
 313             i.print('test')
 314             with i:
 315                 i.print('-ing')
 316                 with i:
 317                     i.print('1, 2, 3')
 318
 319     Yields::
 320
 321         test
 322                 -ing
 323                         1, 2, 3
 324
 325     """
 326
 327     def __init__(
 328         self,
 329         *,
 330         pad_prefix: Optional[str] = None,
 331         pad_char: str = ' ',
 332         pad_count: int = 4,
 333     ):
 334         self.level = -1
 335         if pad_prefix is not None:
 336             self.pad_prefix = pad_prefix
 337         else:
 338             self.pad_prefix = ''
 339         self.padding = pad_char * pad_count
 340
 341     def __enter__(self):
 342         self.level += 1
 343         return self
 344
 345     def __exit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
 346         self.level -= 1
 347         if self.level < -1:
 348             self.level = -1
 349         return False
 350
 351     def print(self, *arg, **kwargs):
 352         text = string_utils.sprintf(*arg, **kwargs)
 353         print(self.pad_prefix + self.padding * self.level + text, end='')
 354
 355
 356 def header(
 357     title: str,
 358     *,
 359     width: Optional[int] = None,
 360     align: Optional[str] = None,
 361     style: Optional[str] = 'solid',
 362     color: Optional[str] = None,
 363 ):
 364     """
 365     Returns a nice header line with a title.
 366
 367     >>> header('title', width=60, style='ascii')
 368     '----[ title ]-----------------------------------------------'
 369
 370     """
 371     if not width:
 372         try:
 373             width = get_console_rows_columns().columns
 374         except Exception:
 375             width = 80
 376     if not align:
 377         align = 'left'
 378     if not style:
 379         style = 'ascii'
 380
 381     text_len = len(string_utils.strip_ansi_sequences(title))
 382     if align == 'left':
 383         left = 4
 384         right = width - (left + text_len + 4)
 385     elif align == 'right':
 386         right = 4
 387         left = width - (right + text_len + 4)
 388     else:
 389         left = int((width - (text_len + 4)) / 2)
 390         right = left
 391         while left + text_len + 4 + right < width:
 392             right += 1
 393
 394     if style == 'solid':
 395         line_char = '━'
 396         begin = ''
 397         end = ''
 398     elif style == 'dashed':
 399         line_char = '┅'
 400         begin = ''
 401         end = ''
 402     else:
 403         line_char = '-'
 404         begin = '['
 405         end = ']'
 406     if color:
 407         col = color
 408         reset_seq = reset()
 409     else:
 410         col = ''
 411         reset_seq = ''
 412     return line_char * left + begin + col + ' ' + title + ' ' + reset_seq + end + line_char * right
 413
 414
 415 def box(
 416     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 417 ) -> str:
 418     assert width > 4
 419     if text is not None:
 420         text = justify_text(text, width=width - 4, alignment='l')
 421     return preformatted_box(title, text, width=width, color=color)
 422
 423
 424 def preformatted_box(
 425     title: Optional[str] = None, text: Optional[str] = None, *, width=80, color: str = ''
 426 ) -> str:
 427     assert width > 4
 428     ret = ''
 429     if color == '':
 430         rset = ''
 431     else:
 432         rset = reset()
 433     w = width - 2
 434     ret += color + '╭' + '─' * w + '╮' + rset + '\n'
 435     if title is not None:
 436         ret += (
 437             color
 438             + '│'
 439             + rset
 440             + justify_string(title, width=w, alignment='c')
 441             + color
 442             + '│'
 443             + rset
 444             + '\n'
 445         )
 446         ret += color + '│' + ' ' * w + '│' + rset + '\n'
 447     if text is not None:
 448         for line in text.split('\n'):
 449             tw = len(string_utils.strip_ansi_sequences(line))
 450             assert tw <= w
 451             ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
 452     ret += color + '╰' + '─' * w + '╯' + rset + '\n'
 453     return ret
 454
 455
 456 def print_box(
 457     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 458 ) -> None:
 459     """Draws a box with nice rounded corners.
 460
 461     >>> print_box('Title', 'This is text', width=30)
 462     ╭────────────────────────────╮
 463     │            Title           │
 464     │                            │
 465     │ This is text               │
 466     ╰────────────────────────────╯
 467
 468     >>> print_box(None, 'OK', width=6)
 469     ╭────╮
 470     │ OK │
 471     ╰────╯
 472
 473     """
 474     print(preformatted_box(title, text, width=width, color=color), end='')
 475
 476
 477 if __name__ == '__main__':
 478     import doctest
 479
 480     doctest.testmod()