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     with Indenter(pad_count = 8) as i:
 311         i.print('test')
 312         with i:
 313             i.print('-ing')
 314             with i:
 315                 i.print('1, 2, 3')
 316
 317     """
 318
 319     def __init__(
 320         self,
 321         *,
 322         pad_prefix: Optional[str] = None,
 323         pad_char: str = ' ',
 324         pad_count: int = 4,
 325     ):
 326         self.level = -1
 327         if pad_prefix is not None:
 328             self.pad_prefix = pad_prefix
 329         else:
 330             self.pad_prefix = ''
 331         self.padding = pad_char * pad_count
 332
 333     def __enter__(self):
 334         self.level += 1
 335         return self
 336
 337     def __exit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
 338         self.level -= 1
 339         if self.level < -1:
 340             self.level = -1
 341         return False
 342
 343     def print(self, *arg, **kwargs):
 344         text = string_utils.sprintf(*arg, **kwargs)
 345         print(self.pad_prefix + self.padding * self.level + text, end='')
 346
 347
 348 def header(
 349     title: str,
 350     *,
 351     width: Optional[int] = None,
 352     align: Optional[str] = None,
 353     style: Optional[str] = 'solid',
 354     color: Optional[str] = None,
 355 ):
 356     """
 357     Returns a nice header line with a title.
 358
 359     >>> header('title', width=60, style='ascii')
 360     '----[ title ]-----------------------------------------------'
 361
 362     """
 363     if not width:
 364         try:
 365             width = get_console_rows_columns().columns
 366         except Exception:
 367             width = 80
 368     if not align:
 369         align = 'left'
 370     if not style:
 371         style = 'ascii'
 372
 373     text_len = len(string_utils.strip_ansi_sequences(title))
 374     if align == 'left':
 375         left = 4
 376         right = width - (left + text_len + 4)
 377     elif align == 'right':
 378         right = 4
 379         left = width - (right + text_len + 4)
 380     else:
 381         left = int((width - (text_len + 4)) / 2)
 382         right = left
 383         while left + text_len + 4 + right < width:
 384             right += 1
 385
 386     if style == 'solid':
 387         line_char = '━'
 388         begin = ''
 389         end = ''
 390     elif style == 'dashed':
 391         line_char = '┅'
 392         begin = ''
 393         end = ''
 394     else:
 395         line_char = '-'
 396         begin = '['
 397         end = ']'
 398     if color:
 399         col = color
 400         reset_seq = reset()
 401     else:
 402         col = ''
 403         reset_seq = ''
 404     return line_char * left + begin + col + ' ' + title + ' ' + reset_seq + end + line_char * right
 405
 406
 407 def box(
 408     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 409 ) -> str:
 410     assert width > 4
 411     if text is not None:
 412         text = justify_text(text, width=width - 4, alignment='l')
 413     return preformatted_box(title, text, width=width, color=color)
 414
 415
 416 def preformatted_box(
 417     title: Optional[str] = None, text: Optional[str] = None, *, width=80, color: str = ''
 418 ) -> str:
 419     assert width > 4
 420     ret = ''
 421     if color == '':
 422         rset = ''
 423     else:
 424         rset = reset()
 425     w = width - 2
 426     ret += color + '╭' + '─' * w + '╮' + rset + '\n'
 427     if title is not None:
 428         ret += (
 429             color
 430             + '│'
 431             + rset
 432             + justify_string(title, width=w, alignment='c')
 433             + color
 434             + '│'
 435             + rset
 436             + '\n'
 437         )
 438         ret += color + '│' + ' ' * w + '│' + rset + '\n'
 439     if text is not None:
 440         for line in text.split('\n'):
 441             tw = len(string_utils.strip_ansi_sequences(line))
 442             assert tw <= w
 443             ret += color + '│ ' + rset + line + ' ' * (w - tw - 2) + color + ' │' + rset + '\n'
 444     ret += color + '╰' + '─' * w + '╯' + rset + '\n'
 445     return ret
 446
 447
 448 def print_box(
 449     title: Optional[str] = None, text: Optional[str] = None, *, width: int = 80, color: str = ''
 450 ) -> None:
 451     """Draws a box with nice rounded corners.
 452
 453     >>> print_box('Title', 'This is text', width=30)
 454     ╭────────────────────────────╮
 455     │            Title           │
 456     │                            │
 457     │ This is text               │
 458     ╰────────────────────────────╯
 459
 460     >>> print_box(None, 'OK', width=6)
 461     ╭────╮
 462     │ OK │
 463     ╰────╯
 464
 465     """
 466     print(preformatted_box(title, text, width=width, color=color), end='')
 467
 468
 469 if __name__ == '__main__':
 470     import doctest
 471
 472     doctest.testmod()