text_utils.py

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