#!/usr/bin/env python3 # -*- coding: utf-8 -*- """Utilities for dealing with "text".""" import contextlib import logging import math import sys from collections import defaultdict from dataclasses import dataclass from typing import Dict, Generator, List, Literal, Optional, Tuple from ansi import fg, reset logger = logging.getLogger(__file__) @dataclass class RowsColumns: """Row + Column""" rows: int = 0 columns: int = 0 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?!') from e 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_seq=reset(), left_end="[", right_end="]", ) -> str: """Returns a string containing a bar graph. >>> bar_graph(0.5, fgcolor='', reset_seq='') '[███████████████████████████████████ ] 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_seq + right_end + " " + text ) def sparkline(numbers: List[float]) -> Tuple[float, float, str]: """ Makes a "sparkline" little inline histogram graph. Auto scales. >>> sparkline([1, 2, 3, 5, 10, 3, 5, 7]) (1, 10, '▁▁▂▄█▂▄▆') >>> sparkline([104, 99, 93, 96, 82, 77, 85, 73]) (73, 104, '█▇▆▆▃▂▄▁') """ _bar = '▁▂▃▄▅▆▇█' # Unicode: 9601, 9602, 9603, 9604, 9605, 9606, 9607, 9608 barcount = len(_bar) min_num, max_num = min(numbers), max(numbers) span = max_num - min_num sline = ''.join( _bar[min([barcount - 1, int((n - min_num) / span * barcount)])] for n in numbers ) return min_num, max_num, sline 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(contextlib.AbstractContextManager): """ 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) -> Literal[False]: self.level -= 1 if self.level < -1: self.level = -1 return False 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()