text_utils.py

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