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