text_utils.py

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