222 lines
7.1 KiB
Python
222 lines
7.1 KiB
Python
import difflib
|
|
import os
|
|
import sys
|
|
import textwrap
|
|
from typing import Any, Optional, Tuple, Union
|
|
|
|
STDOUT_ENCODING = sys.stdout.encoding if hasattr(sys.stdout, "encoding") else None
|
|
ENCODING = STDOUT_ENCODING or "ascii"
|
|
NO_UTF8 = ENCODING.lower() not in ("utf8", "utf-8")
|
|
|
|
|
|
# Environment variables
|
|
ENV_ANSI_DISABLED = "ANSI_COLORS_DISABLED" # no colors
|
|
|
|
|
|
class MESSAGES(object):
|
|
GOOD = "good"
|
|
FAIL = "fail"
|
|
WARN = "warn"
|
|
INFO = "info"
|
|
|
|
|
|
COLORS = {
|
|
MESSAGES.GOOD: 2,
|
|
MESSAGES.FAIL: 1,
|
|
MESSAGES.WARN: 3,
|
|
MESSAGES.INFO: 4,
|
|
"red": 1,
|
|
"green": 2,
|
|
"yellow": 3,
|
|
"blue": 4,
|
|
"pink": 5,
|
|
"cyan": 6,
|
|
"white": 7,
|
|
"grey": 8,
|
|
"black": 16,
|
|
}
|
|
|
|
|
|
ICONS = {
|
|
MESSAGES.GOOD: "\u2714" if not NO_UTF8 else "[+]",
|
|
MESSAGES.FAIL: "\u2718" if not NO_UTF8 else "[x]",
|
|
MESSAGES.WARN: "\u26a0" if not NO_UTF8 else "[!]",
|
|
MESSAGES.INFO: "\u2139" if not NO_UTF8 else "[i]",
|
|
}
|
|
|
|
INSERT_SYMBOL = "+"
|
|
DELETE_SYMBOL = "-"
|
|
|
|
|
|
def color(
|
|
text: str,
|
|
fg: Optional[Union[str, int]] = None,
|
|
bg: Optional[Union[str, int]] = None,
|
|
bold: bool = False,
|
|
underline: bool = False,
|
|
) -> str:
|
|
"""Color text by applying ANSI escape sequence.
|
|
|
|
text (str): The text to be formatted.
|
|
fg (Optional[Union[str, int]]): Optional foreground color. String name or 0 - 256 (see COLORS).
|
|
bg (Optional[Union[str, int]]): Optional background color. String name or 0 - 256 (see COLORS).
|
|
bold (bool): Format text in bold.
|
|
underline (bool): Underline text.
|
|
RETURNS (str): The formatted text.
|
|
"""
|
|
fg = COLORS.get(fg, fg) # type: ignore
|
|
bg = COLORS.get(bg, bg) # type: ignore
|
|
if not any([fg, bg, bold]):
|
|
return text
|
|
styles = []
|
|
if bold:
|
|
styles.append("1")
|
|
if underline:
|
|
styles.append("4")
|
|
if fg:
|
|
styles.append("38;5;{}".format(fg))
|
|
if bg:
|
|
styles.append("48;5;{}".format(bg))
|
|
return "\x1b[{}m{}\x1b[0m".format(";".join(styles), text)
|
|
|
|
|
|
def wrap(text: Any, wrap_max: int = 80, indent: int = 4) -> str:
|
|
"""Wrap text at given width using textwrap module.
|
|
|
|
text (Any): The text to wrap.
|
|
wrap_max (int): Maximum line width, including indentation. Defaults to 80.
|
|
indent (int): Number of spaces used for indentation. Defaults to 4.
|
|
RETURNS (str): The wrapped text with line breaks.
|
|
"""
|
|
indent_str = indent * " "
|
|
wrap_width = wrap_max - len(indent_str)
|
|
text = str(text)
|
|
return textwrap.fill(
|
|
text,
|
|
width=wrap_width,
|
|
initial_indent=indent_str,
|
|
subsequent_indent=indent_str,
|
|
break_long_words=False,
|
|
break_on_hyphens=False,
|
|
)
|
|
|
|
|
|
def format_repr(obj: Any, max_len: int = 50, ellipsis: str = "...") -> str:
|
|
"""Wrapper around `repr()` to print shortened and formatted string version.
|
|
|
|
obj: The object to represent.
|
|
max_len (int): Maximum string length. Longer strings will be cut in the
|
|
middle so only the beginning and end is displayed, separated by ellipsis.
|
|
ellipsis (str): Ellipsis character(s), e.g. "...".
|
|
RETURNS (str): The formatted representation.
|
|
"""
|
|
string = repr(obj)
|
|
if len(string) >= max_len:
|
|
half = int(max_len / 2)
|
|
return "{} {} {}".format(string[:half], ellipsis, string[-half:])
|
|
else:
|
|
return string
|
|
|
|
|
|
def diff_strings(
|
|
a: str,
|
|
b: str,
|
|
fg: Union[str, int] = "black",
|
|
bg: Union[Tuple[str, str], Tuple[int, int]] = ("green", "red"),
|
|
add_symbols: bool = False,
|
|
) -> str:
|
|
"""Compare two strings and return a colored diff with red/green background
|
|
for deletion and insertions.
|
|
|
|
a (str): The first string to diff.
|
|
b (str): The second string to diff.
|
|
fg (Union[str, int]): Foreground color. String name or 0 - 256 (see COLORS).
|
|
bg (Union[Tuple[str, str], Tuple[int, int]]): Background colors as
|
|
(insert, delete) tuple of string name or 0 - 256 (see COLORS).
|
|
add_symbols (bool): Whether to add symbols before the diff lines. Uses '+'
|
|
for inserts and '-' for deletions. Default is False.
|
|
RETURNS (str): The formatted diff.
|
|
"""
|
|
a_list = a.split("\n")
|
|
b_list = b.split("\n")
|
|
output = []
|
|
matcher = difflib.SequenceMatcher(None, a_list, b_list)
|
|
for opcode, a0, a1, b0, b1 in matcher.get_opcodes():
|
|
if opcode == "equal":
|
|
for item in a_list[a0:a1]:
|
|
output.append(item)
|
|
if opcode == "insert" or opcode == "replace":
|
|
for item in b_list[b0:b1]:
|
|
item = "{} {}".format(INSERT_SYMBOL, item) if add_symbols else item
|
|
output.append(color(item, fg=fg, bg=bg[0]))
|
|
if opcode == "delete" or opcode == "replace":
|
|
for item in a_list[a0:a1]:
|
|
item = "{} {}".format(DELETE_SYMBOL, item) if add_symbols else item
|
|
output.append(color(item, fg=fg, bg=bg[1]))
|
|
return "\n".join(output)
|
|
|
|
|
|
def get_raw_input(
|
|
description: str, default: Optional[Union[str, bool]] = False, indent: int = 4
|
|
) -> str:
|
|
"""Get user input from the command line via raw_input / input.
|
|
|
|
description (str): Text to display before prompt.
|
|
default (Optional[Union[str, bool]]): Optional default value to display with prompt.
|
|
indent (int): Indentation in spaces.
|
|
RETURNS (str): User input.
|
|
"""
|
|
additional = " (default: {})".format(default) if default else ""
|
|
prompt = wrap("{}{}: ".format(description, additional), indent=indent)
|
|
user_input = input(prompt)
|
|
return user_input
|
|
|
|
|
|
def locale_escape(string: Any, errors: str = "replace") -> str:
|
|
"""Mangle non-supported characters, for savages with ASCII terminals.
|
|
|
|
string (Any): The string to escape.
|
|
errors (str): The str.encode errors setting. Defaults to `"replace"`.
|
|
RETURNS (str): The escaped string.
|
|
"""
|
|
string = str(string)
|
|
string = string.encode(ENCODING, errors).decode("utf8")
|
|
return string
|
|
|
|
|
|
def can_render(string: str) -> bool:
|
|
"""Check if terminal can render unicode characters, e.g. special loading
|
|
icons. Can be used to display fallbacks for ASCII terminals.
|
|
|
|
string (str): The string to render.
|
|
RETURNS (bool): Whether the terminal can render the text.
|
|
"""
|
|
try:
|
|
string.encode(ENCODING)
|
|
return True
|
|
except UnicodeEncodeError:
|
|
return False
|
|
|
|
|
|
def supports_ansi() -> bool:
|
|
"""Returns True if the running system's terminal supports ANSI escape
|
|
sequences for color, formatting etc. and False otherwise.
|
|
|
|
RETURNS (bool): Whether the terminal supports ANSI colors.
|
|
"""
|
|
if os.getenv(ENV_ANSI_DISABLED):
|
|
return False
|
|
# We require colorama on Windows Python 3.7+, but we might be running on Unix, or we
|
|
# might be running on Windows Python 3.6. In both cases, colorama might be missing,
|
|
# *or* there might by accident happen to be an install of an old version that
|
|
# doesn't have just_fix_windows_console. So we need to confirm not just that we can
|
|
# import colorama, but that we can import just_fix_windows_console.
|
|
try:
|
|
from colorama import just_fix_windows_console
|
|
except ImportError:
|
|
if sys.platform == "win32" and "ANSICON" not in os.environ:
|
|
return False
|
|
else:
|
|
just_fix_windows_console()
|
|
return True
|