98 lines
2.8 KiB
Python
98 lines
2.8 KiB
Python
|
import os
|
||
|
|
import sys
|
||
|
|
|
||
|
|
if sys.platform == 'win32': # pragma: no cover (windows)
|
||
|
|
def _enable() -> None:
|
||
|
|
from ctypes import POINTER
|
||
|
|
from ctypes import windll
|
||
|
|
from ctypes import WinError
|
||
|
|
from ctypes import WINFUNCTYPE
|
||
|
|
from ctypes.wintypes import BOOL
|
||
|
|
from ctypes.wintypes import DWORD
|
||
|
|
from ctypes.wintypes import HANDLE
|
||
|
|
|
||
|
|
STD_OUTPUT_HANDLE = -11
|
||
|
|
ENABLE_VIRTUAL_TERMINAL_PROCESSING = 4
|
||
|
|
|
||
|
|
def bool_errcheck(result, func, args):
|
||
|
|
if not result:
|
||
|
|
raise WinError()
|
||
|
|
return args
|
||
|
|
|
||
|
|
GetStdHandle = WINFUNCTYPE(HANDLE, DWORD)(
|
||
|
|
('GetStdHandle', windll.kernel32), ((1, 'nStdHandle'),),
|
||
|
|
)
|
||
|
|
|
||
|
|
GetConsoleMode = WINFUNCTYPE(BOOL, HANDLE, POINTER(DWORD))(
|
||
|
|
('GetConsoleMode', windll.kernel32),
|
||
|
|
((1, 'hConsoleHandle'), (2, 'lpMode')),
|
||
|
|
)
|
||
|
|
GetConsoleMode.errcheck = bool_errcheck
|
||
|
|
|
||
|
|
SetConsoleMode = WINFUNCTYPE(BOOL, HANDLE, DWORD)(
|
||
|
|
('SetConsoleMode', windll.kernel32),
|
||
|
|
((1, 'hConsoleHandle'), (1, 'dwMode')),
|
||
|
|
)
|
||
|
|
SetConsoleMode.errcheck = bool_errcheck
|
||
|
|
|
||
|
|
# As of Windows 10, the Windows console supports (some) ANSI escape
|
||
|
|
# sequences, but it needs to be enabled using `SetConsoleMode` first.
|
||
|
|
#
|
||
|
|
# More info on the escape sequences supported:
|
||
|
|
# https://msdn.microsoft.com/en-us/library/windows/desktop/mt638032(v=vs.85).aspx
|
||
|
|
stdout = GetStdHandle(STD_OUTPUT_HANDLE)
|
||
|
|
flags = GetConsoleMode(stdout)
|
||
|
|
SetConsoleMode(stdout, flags | ENABLE_VIRTUAL_TERMINAL_PROCESSING)
|
||
|
|
|
||
|
|
try:
|
||
|
|
_enable()
|
||
|
|
except OSError:
|
||
|
|
terminal_supports_color = False
|
||
|
|
else:
|
||
|
|
terminal_supports_color = True
|
||
|
|
else: # pragma: win32 no cover
|
||
|
|
terminal_supports_color = True
|
||
|
|
|
||
|
|
RED = '\033[41m'
|
||
|
|
GREEN = '\033[42m'
|
||
|
|
YELLOW = '\033[43;30m'
|
||
|
|
TURQUOISE = '\033[46;30m'
|
||
|
|
SUBTLE = '\033[2m'
|
||
|
|
NORMAL = '\033[m'
|
||
|
|
|
||
|
|
|
||
|
|
def format_color(text: str, color: str, use_color_setting: bool) -> str:
|
||
|
|
"""Format text with color.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
text - Text to be formatted with color if `use_color`
|
||
|
|
color - The color start string
|
||
|
|
use_color_setting - Whether or not to color
|
||
|
|
"""
|
||
|
|
if use_color_setting:
|
||
|
|
return f'{color}{text}{NORMAL}'
|
||
|
|
else:
|
||
|
|
return text
|
||
|
|
|
||
|
|
|
||
|
|
COLOR_CHOICES = ('auto', 'always', 'never')
|
||
|
|
|
||
|
|
|
||
|
|
def use_color(setting: str) -> bool:
|
||
|
|
"""Choose whether to use color based on the command argument.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
setting - Either `auto`, `always`, or `never`
|
||
|
|
"""
|
||
|
|
if setting not in COLOR_CHOICES:
|
||
|
|
raise ValueError(setting)
|
||
|
|
|
||
|
|
return (
|
||
|
|
setting == 'always' or (
|
||
|
|
setting == 'auto' and
|
||
|
|
sys.stdout.isatty() and
|
||
|
|
terminal_supports_color and
|
||
|
|
os.getenv('TERM') != 'dumb'
|
||
|
|
)
|
||
|
|
)
|