#!/usr/bin/env python3 """Utilities for dealing with "text".""" from collections import defaultdict import logging import math import sys from typing import Dict, Generator, List, NamedTuple, Optional from ansi import fg, reset logger = logging.getLogger(__file__) class RowsColumns(NamedTuple): rows: int columns: int def get_console_rows_columns() -> RowsColumns: """Returns the number of rows/columns on the current console.""" from exec_utils import cmd try: rows, columns = cmd( "stty size", timeout_seconds=1.0, ).split() except Exception as e: logger.exception(e) raise Exception('Can\'t determine console size?!') return RowsColumns(int(rows), int(columns)) def progress_graph( current: int, total: int, *, width=70, fgcolor=fg("school bus yellow"), left_end="[", right_end="]", redraw=True, ) -> None: """Draws a progress graph.""" percent = current / total ret = "\r" if redraw else "\n" bar = bar_graph( percent, include_text=True, width=width, fgcolor=fgcolor, left_end=left_end, right_end=right_end, ) print(bar, end=ret, flush=True, file=sys.stderr) def bar_graph( percentage: float, *, include_text=True, width=70, fgcolor=fg("school bus yellow"), reset=reset(), left_end="[", right_end="]", ) -> str: """Returns a string containing a bar graph. >>> bar_graph(0.5, fgcolor='', reset='') '[███████████████████████████████████ ] 50.0%' """ if percentage < 0.0 or percentage > 1.0: raise ValueError(percentage) if include_text: text = f"{percentage*100.0:2.1f}%" else: text = "" whole_width = math.floor(percentage * width) if whole_width == width: whole_width -= 1 part_char = "▉" elif whole_width == 0 and percentage > 0.0: part_char = "▏" else: remainder_width = (percentage * width) % 1 part_width = math.floor(remainder_width * 8) part_char = [" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉"][part_width] return ( left_end + fgcolor + "█" * whole_width + part_char + " " * (width - whole_width - 1) + reset + right_end + " " + text ) def distribute_strings( strings: List[str], *, width: int = 80, alignment: str = "c", padding: str = " ", ) -> str: """ Distributes strings into a line with a particular justification. >>> distribute_strings(['this', 'is', 'a', 'test'], width=40) ' this is a test ' >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='l') 'this is a test ' >>> distribute_strings(['this', 'is', 'a', 'test'], width=40, alignment='r') ' this is a test' """ subwidth = math.floor(width / len(strings)) retval = "" for string in strings: string = justify_string( string, width=subwidth, alignment=alignment, padding=padding ) retval += string while len(retval) > width: retval = retval.replace(' ', ' ', 1) while len(retval) < width: retval = retval.replace(' ', ' ', 1) return retval def justify_string_by_chunk(string: str, width: int = 80, padding: str = " ") -> str: """ Justifies a string. >>> justify_string_by_chunk("This is a test", 40) 'This is a test' >>> justify_string_by_chunk("This is a test", 20) 'This is a test' """ padding = padding[0] first, *rest, last = string.split() w = width - (len(first) + 1 + len(last) + 1) ret = ( first + padding + distribute_strings(rest, width=w, padding=padding) + padding + last ) return ret def justify_string( string: str, *, width: int = 80, alignment: str = "c", padding: str = " " ) -> str: """Justify a string. >>> justify_string('This is another test', width=40, alignment='c') ' This is another test ' >>> justify_string('This is another test', width=40, alignment='l') 'This is another test ' >>> justify_string('This is another test', width=40, alignment='r') ' This is another test' >>> justify_string('This is another test', width=40, alignment='j') 'This is another test' """ alignment = alignment[0] padding = padding[0] while len(string) < width: if alignment == "l": string += padding elif alignment == "r": string = padding + string elif alignment == "j": return justify_string_by_chunk(string, width=width, padding=padding) elif alignment == "c": if len(string) % 2 == 0: string += padding else: string = padding + string else: raise ValueError return string def justify_text(text: str, *, width: int = 80, alignment: str = "c") -> str: """ Justifies text. >>> justify_text('This is a test of the emergency broadcast system. This is only a test.', ... width=40, alignment='j') #doctest: +NORMALIZE_WHITESPACE 'This is a test of the emergency\\nbroadcast system. This is only a test.' """ retval = "" line = "" for word in text.split(): if len(line) + len(word) > width: line = line[1:] line = justify_string(line, width=width, alignment=alignment) retval = retval + "\n" + line line = "" line = line + " " + word if len(line) > 0: retval += "\n" + line[1:] return retval[1:] def generate_padded_columns(text: List[str]) -> Generator: max_width: Dict[int, int] = defaultdict(int) for line in text: for pos, word in enumerate(line.split()): max_width[pos] = max(max_width[pos], len(word)) for line in text: out = "" for pos, word in enumerate(line.split()): width = max_width[pos] word = justify_string(word, width=width, alignment='l') out += f'{word} ' yield out def wrap_string(text: str, n: int) -> str: chunks = text.split() out = '' width = 0 for chunk in chunks: if width + len(chunk) > n: out += '\n' width = 0 out += chunk + ' ' width += len(chunk) + 1 return out class Indenter(object): """ with Indenter(pad_count = 8) as i: i.print('test') with i: i.print('-ing') with i: i.print('1, 2, 3') """ def __init__( self, *, pad_prefix: Optional[str] = None, pad_char: str = ' ', pad_count: int = 4, ): self.level = -1 if pad_prefix is not None: self.pad_prefix = pad_prefix else: self.pad_prefix = '' self.padding = pad_char * pad_count def __enter__(self): self.level += 1 return self def __exit__(self, exc_type, exc_value, exc_tb): self.level -= 1 if self.level < -1: self.level = -1 def print(self, *arg, **kwargs): import string_utils text = string_utils.sprintf(*arg, **kwargs) print(self.pad_prefix + self.padding * self.level + text, end='') def header(title: str, *, width: int = 80, color: str = ''): """ Returns a nice header line with a title. >>> header('title', width=60, color='') '----[ title ]-----------------------------------------------' """ w = width w -= len(title) + 4 if w >= 4: left = 4 * '-' right = (w - 4) * '-' if color != '' and color is not None: r = reset() else: r = '' return f'{left}[ {color}{title}{r} ]{right}' else: return '' if __name__ == '__main__': import doctest doctest.testmod()