72 lines
1.8 KiB
Python
72 lines
1.8 KiB
Python
"""WizardKit: ANSI control/escape functions"""
|
|
# vim: sts=2 sw=2 ts=2
|
|
|
|
import itertools
|
|
import logging
|
|
import pathlib
|
|
|
|
# STATIC VARIABLES
|
|
LOG = logging.getLogger(__name__)
|
|
COLORS = {
|
|
'CLEAR': '\033[0m',
|
|
'RED': '\033[31m',
|
|
'RED_BLINK': '\033[31;5m',
|
|
'ORANGE': '\033[31;1m',
|
|
'ORANGE_RED': '\033[1;31;41m',
|
|
'GREEN': '\033[32m',
|
|
'YELLOW': '\033[33m',
|
|
'YELLOW_BLINK': '\033[33;5m',
|
|
'BLUE': '\033[34m',
|
|
'PURPLE': '\033[35m',
|
|
'CYAN': '\033[36m',
|
|
}
|
|
|
|
|
|
# Functions
|
|
def clear_screen() -> None:
|
|
"""Clear screen using ANSI escape."""
|
|
print('\033c', end='', flush=True)
|
|
|
|
|
|
def color_string(strings, colors, sep=' ') -> str:
|
|
"""Build colored string using ANSI escapes, returns str."""
|
|
clear_code = COLORS['CLEAR']
|
|
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),)
|
|
|
|
# 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}')
|
|
|
|
# Done
|
|
return sep.join(msg)
|
|
|
|
|
|
def strip_colors(string) -> str:
|
|
"""Strip known ANSI color escapes from string, returns str."""
|
|
LOG.debug('string: %s', string)
|
|
for color in COLORS.values():
|
|
string = string.replace(color, '')
|
|
return string
|
|
|
|
|
|
if __name__ == '__main__':
|
|
print("This file is not meant to be called directly.")
|