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