diff --git a/scripts/wk/ui/ansi.py b/scripts/wk/ui/ansi.py index c31b6459..b65fe2c9 100644 --- a/scripts/wk/ui/ansi.py +++ b/scripts/wk/ui/ansi.py @@ -3,7 +3,8 @@ import itertools import logging -import pathlib + +from typing import Iterable # STATIC VARIABLES LOG = logging.getLogger(__name__) @@ -28,39 +29,36 @@ def clear_screen() -> None: print('\033c', end='', flush=True) -def color_string(strings, colors, sep=' ') -> str: +def color_string( + strings: Iterable[str] | str, + colors: Iterable[str | None] | str, + sep=' ', + ) -> str: """Build colored string using ANSI escapes, returns str.""" - clear_code = COLORS['CLEAR'] + data = {'strings': strings, 'colors': colors} msg = [] - # Convert to tuples if necessary - if isinstance(strings, (str, pathlib.Path)): - strings = (strings,) - if isinstance(colors, (str, pathlib.Path)): - colors = (colors,) - - # Convert to strings if necessary - try: - iter(strings) - except TypeError: - # Assuming single element passed, convert to string - strings = (str(strings),) - try: - iter(colors) - except TypeError: - # Assuming single element passed, convert to string - colors = (str(colors),) + # Convert input to tuples of strings + for k, v in data.items(): + if isinstance(v, str): + # Avoid splitting string into a list of characters + data[k] = (v,) + try: + iter(v) + except TypeError: + # Assuming single element passed, convert to string + data[k] = (str(v),) # Build new string with color escapes added - for string, color in itertools.zip_longest(strings, colors): - color_code = COLORS.get(color, clear_code) - msg.append(f'{color_code}{string}{clear_code}') + for string, color in itertools.zip_longest(data['strings'], data['colors']): + color_code = COLORS.get(str(color), COLORS['CLEAR']) + msg.append(f'{color_code}{string}{COLORS["CLEAR"]}') # Done return sep.join(msg) -def strip_colors(string) -> str: +def strip_colors(string: str) -> str: """Strip known ANSI color escapes from string, returns str.""" LOG.debug('string: %s', string) for color in COLORS.values():