text_utils.py

   1 #!/usr/bin/env python3
   2
   3 """Utilities for dealing with "text"."""
   4
   5 import logging
   6 import math
   7 import sys
   8 from collections import defaultdict
   9 from typing import Dict, Generator, List, NamedTuple, Optional
  10
  11 from ansi import fg, reset
  12
  13 logger = logging.getLogger(__file__)
  14
  15
  16 class RowsColumns(NamedTuple):
  17     rows: int
  18     columns: int
  19
  20
  21 def get_console_rows_columns() -> RowsColumns:
  22     """Returns the number of rows/columns on the current console."""
  23
  24     from exec_utils import cmd
  25
  26     try:
  27         rows, columns = cmd(
  28             "stty size",
  29             timeout_seconds=1.0,
  30         ).split()
  31     except Exception as e:
  32         logger.exception(e)
  33         raise Exception('Can\'t determine console size?!')
  34     return RowsColumns(int(rows), int(columns))
  35
  36
  37 def progress_graph(
  38     current: int,
  39     total: int,
  40     *,
  41     width=70,
  42     fgcolor=fg("school bus yellow"),
  43     left_end="[",
  44     right_end="]",
  45     redraw=True,
  46 ) -> None:
  47     """Draws a progress graph."""
  48
  49     percent = current / total
  50     ret = "\r" if redraw else "\n"
  51     bar = bar_graph(
  52         percent,
  53         include_text=True,
  54         width=width,
  55         fgcolor=fgcolor,
  56         left_end=left_end,
  57         right_end=right_end,
  58     )
  59     print(bar, end=ret, flush=True, file=sys.stderr)
  60
  61
  62 def bar_graph(
  63     percentage: float,
  64     *,
  65     include_text=True,
  66     width=70,
  67     fgcolor=fg("school bus yellow"),
  68     reset=reset(),
  69     left_end="[",
  70     right_end="]",
  71 ) -> str:
  72     """Returns a string containing a bar graph.
  73
  74     >>> bar_graph(0.5, fgcolor='', reset='')
  75     '[███████████████████████████████████                                   ] 50.0%'
  76
  77     """
  78
  79     if percentage < 0.0 or percentage > 1.0:
  80         raise ValueError(percentage)
  81     if include_text:
  82         text = f"{percentage*100.0:2.1f}%"
  83     else:
  84         text = ""
  85     whole_width = math.floor(percentage * width)
  86     if whole_width == width:
  87         whole_width -= 1
  88         part_char = "▉"
  89     elif whole_width == 0 and percentage > 0.0:
  90         part_char = "▏"
  91     else:
  92         remainder_width = (percentage * width) % 1
  93         part_width = math.floor(remainder_width * 8)
  94         part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width]
  95     return (
  96         left_end
  97         + fgcolor
  98         + "█" * whole_width
  99         + part_char
 100         + " " * (width - whole_width - 1)
 101         + reset
 102         + right_end
 103         + " "
 104         + text
 105     )
 106
 107
 108 def distribute_strings(
 109     strings: List[str],
 110     *,
 111     width: int = 80,
 112     alignment: str = "c",
 113     padding: str = " ",
 114 ) -> str:
 115     """
 116     Distributes strings into a line with a particular justification.
 117
 118     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40)
 119     '   this       is         a       test   '
 120     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='l')
 121     'this      is        a         test      '
 122     >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='r')
 123     '      this        is         a      test'
 124
 125     """
 126     subwidth = math.floor(width / len(strings))
 127     retval = ""
 128     for string in strings:
 129         string = justify_string(
 130             string, width=subwidth, alignment=alignment, padding=padding
 131         )
 132         retval += string
 133     while len(retval) > width:
 134         retval = retval.replace('  ', ' ', 1)
 135     while len(retval) < width:
 136         retval = retval.replace(' ', '  ', 1)
 137     return retval
 138
 139
 140 def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str:
 141     """
 142     Justifies a string.
 143
 144     >>> justify_string_by_chunk("This is a test", 40)
 145     'This       is              a        test'
 146     >>> justify_string_by_chunk("This is a test", 20)
 147     'This  is    a   test'
 148
 149     """
 150     padding = padding[0]
 151     first, *rest, last = string.split()
 152     w = width - (len(first) + 1 + len(last) + 1)
 153     ret = (
 154         first
 155         + padding
 156         + distribute_strings(rest, width=w, padding=padding)
 157         + padding
 158         + last
 159     )
 160     return ret
 161
 162
 163 def justify_string(
 164     string: str, *, width: int = 80, alignment: str = "c", padding: str = " "
 165 ) -> str:
 166     """Justify a string.
 167
 168     >>> justify_string('This is another test', width=40, alignment='c')
 169     '          This is another test          '
 170     >>> justify_string('This is another test', width=40, alignment='l')
 171     'This is another test                    '
 172     >>> justify_string('This is another test', width=40, alignment='r')
 173     '                    This is another test'
 174     >>> justify_string('This is another test', width=40, alignment='j')
 175     'This       is           another     test'
 176
 177     """
 178     alignment = alignment[0]
 179     padding = padding[0]
 180     while len(string) < width:
 181         if alignment == "l":
 182             string += padding
 183         elif alignment == "r":
 184             string = padding + string
 185         elif alignment == "j":
 186             return justify_string_by_chunk(string, width=width, padding=padding)
 187         elif alignment == "c":
 188             if len(string) % 2 == 0:
 189                 string += padding
 190             else:
 191                 string = padding + string
 192         else:
 193             raise ValueError
 194     return string
 195
 196
 197 def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str:
 198     """
 199     Justifies text.
 200
 201     >>> justify_text('This is a test of the emergency broadcast system.  This is only a test.',
 202     ...              width=40, alignment='j')  #doctest: +NORMALIZE_WHITESPACE
 203     'This  is    a  test  of   the  emergency\\nbroadcast system. This is only a test.'
 204     """
 205     retval = ""
 206     line = ""
 207     for word in text.split():
 208         if len(line) + len(word) > width:
 209             line = line[1:]
 210             line = justify_string(line, width=width, alignment=alignment)
 211             retval = retval + "\n" + line
 212             line = ""
 213         line = line + " " + word
 214     if len(line) > 0:
 215         retval += "\n" + line[1:]
 216     return retval[1:]
 217
 218
 219 def generate_padded_columns(text: List[str]) -> Generator:
 220     max_width: Dict[int, int] = defaultdict(int)
 221     for line in text:
 222         for pos, word in enumerate(line.split()):
 223             max_width[pos] = max(max_width[pos], len(word))
 224
 225     for line in text:
 226         out = ""
 227         for pos, word in enumerate(line.split()):
 228             width = max_width[pos]
 229             word = justify_string(word, width=width, alignment='l')
 230             out += f'{word} '
 231         yield out
 232
 233
 234 def wrap_string(text: str, n: int) -> str:
 235     chunks = text.split()
 236     out = ''
 237     width = 0
 238     for chunk in chunks:
 239         if width + len(chunk) > n:
 240             out += '\n'
 241             width = 0
 242         out += chunk + ' '
 243         width += len(chunk) + 1
 244     return out
 245
 246
 247 class Indenter(object):
 248     """
 249     with Indenter(pad_count = 8) as i:
 250         i.print('test')
 251         with i:
 252             i.print('-ing')
 253             with i:
 254                 i.print('1, 2, 3')
 255     """
 256
 257     def __init__(
 258         self,
 259         *,
 260         pad_prefix: Optional[str] = None,
 261         pad_char: str = ' ',
 262         pad_count: int = 4,
 263     ):
 264         self.level = -1
 265         if pad_prefix is not None:
 266             self.pad_prefix = pad_prefix
 267         else:
 268             self.pad_prefix = ''
 269         self.padding = pad_char * pad_count
 270
 271     def __enter__(self):
 272         self.level += 1
 273         return self
 274
 275     def __exit__(self, exc_type, exc_value, exc_tb):
 276         self.level -= 1
 277         if self.level < -1:
 278             self.level = -1
 279
 280     def print(self, *arg, **kwargs):
 281         import string_utils
 282
 283         text = string_utils.sprintf(*arg, **kwargs)
 284         print(self.pad_prefix + self.padding * self.level + text, end='')
 285
 286
 287 def header(title: str, *, width: int = 80, color: str = ''):
 288     """
 289     Returns a nice header line with a title.
 290
 291     >>> header('title', width=60, color='')
 292     '----[ title ]-----------------------------------------------'
 293
 294     """
 295     w = width
 296     w -= len(title) + 4
 297     if w >= 4:
 298         left = 4 * '-'
 299         right = (w - 4) * '-'
 300         if color != '' and color is not None:
 301             r = reset()
 302         else:
 303             r = ''
 304         return f'{left}[ {color}{title}{r} ]{right}'
 305     else:
 306         return ''
 307
 308
 309 if __name__ == '__main__':
 310     import doctest
 311
 312     doctest.testmod()