diff --git a/knack/_win_vt.py b/knack/_win_vt.py new file mode 100644 index 0000000..0cbdd2c --- /dev/null +++ b/knack/_win_vt.py @@ -0,0 +1,77 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +""" +Enable Virtual Terminal mode on legacy Windows terminal to support ANSI escape sequences. +Migrated from https://github.com/Azure/azure-cli/pull/12942 +""" + +from ctypes import WinDLL, get_last_error, byref +from ctypes.wintypes import HANDLE, LPDWORD, DWORD +from msvcrt import get_osfhandle # pylint: disable=import-error +from knack.log import get_logger + +logger = get_logger(__name__) + +ERROR_INVALID_PARAMETER = 0x0057 +ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004 + + +def _check_zero(result, _, args): + if not result: + raise OSError(get_last_error()) + return args + + +# See: +# - https://docs.microsoft.com/en-us/windows/console/getconsolemode +# - https://docs.microsoft.com/en-us/windows/console/setconsolemode +kernel32 = WinDLL("kernel32", use_last_error=True) +kernel32.GetConsoleMode.errcheck = _check_zero +kernel32.GetConsoleMode.argtypes = (HANDLE, LPDWORD) +kernel32.SetConsoleMode.errcheck = _check_zero +kernel32.SetConsoleMode.argtypes = (HANDLE, DWORD) + + +def _get_conout_mode(): + with open("CONOUT$", "w") as conout: # pylint: disable=unspecified-encoding + mode = DWORD() + conout_handle = get_osfhandle(conout.fileno()) + kernel32.GetConsoleMode(conout_handle, byref(mode)) + return mode.value + + +def _set_conout_mode(mode): + with open("CONOUT$", "w") as conout: # pylint: disable=unspecified-encoding + conout_handle = get_osfhandle(conout.fileno()) + kernel32.SetConsoleMode(conout_handle, mode) + + +def _update_conout_mode(mode): + old_mode = _get_conout_mode() + if old_mode & mode != mode: + mode = old_mode | mode # pylint: disable=unsupported-binary-operation + _set_conout_mode(mode) + + +def enable_vt_mode(): + """Enables virtual terminal mode for Windows 10 console. + + Windows 10 supports VT (virtual terminal) / ANSI escape sequences since version 1607. + + cmd.exe enables VT mode, but only for itself. It disables VT mode before starting other programs, + and also at shutdown (See: https://bugs.python.org/issue30075). + + Return True if success, else False. + """ + try: + _update_conout_mode(ENABLE_VIRTUAL_TERMINAL_PROCESSING) + return True + except OSError as e: + if e.errno == ERROR_INVALID_PARAMETER: + logger.debug("Unable to enable virtual terminal processing for legacy Windows terminal.") + else: + logger.debug("Unable to enable virtual terminal processing: %s.", e.errno) + return False diff --git a/knack/cli.py b/knack/cli.py index dda8a44..8178be7 100644 --- a/knack/cli.py +++ b/knack/cli.py @@ -97,8 +97,8 @@ def __init__(self, self.only_show_errors = self.config.getboolean('core', 'only_show_errors', fallback=False) self.enable_color = self._should_enable_color() - # Init colorama only in Windows legacy terminal - self._should_init_colorama = self.enable_color and sys.platform == 'win32' and not is_modern_terminal() + # Enable VT mode only in Windows legacy terminal + self._should_enable_vt_mode = self.enable_color and sys.platform == 'win32' and not is_modern_terminal() @staticmethod def _should_show_version(args): @@ -205,12 +205,14 @@ def invoke(self, args, initial_invocation_data=None, out_file=None): exit_code = 0 try: out_file = out_file or self.out_file - if out_file is sys.stdout and self._should_init_colorama: - self.init_debug_log.append("Init colorama.") - import colorama - colorama.init() - # point out_file to the new sys.stdout which is overwritten by colorama - out_file = sys.stdout + + # Enable VT mode if necessary + if out_file is sys.stdout and self._should_enable_vt_mode: + self.init_debug_log.append("Enable VT mode.") + from ._win_vt import enable_vt_mode + if not enable_vt_mode(): + # Disable color if we can't enable it + self.enable_color = False args = self.completion.get_completion_args() or args @@ -249,10 +251,6 @@ def invoke(self, args, initial_invocation_data=None, out_file=None): finally: self.raise_event(EVENT_CLI_POST_EXECUTE) - if self._should_init_colorama: - import colorama - colorama.deinit() - return exit_code def _should_enable_color(self): @@ -262,7 +260,6 @@ def _should_enable_color(self): # - Otherwise, if the downstream command doesn't support color, Knack will fail with # BrokenPipeError: [Errno 32] Broken pipe, like `az --version | head --lines=1` # https://github.com/Azure/azure-cli/issues/13413 - # - May also hit https://github.com/tartley/colorama/issues/200 # 3. stderr is a tty. # - Otherwise, the output in stderr won't have LEVEL tag # 4. out_file is stdout diff --git a/requirements.txt b/requirements.txt index 39172ca..2a9ddfe 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ argcomplete==1.12.2 -colorama==0.4.4 flake8==4.0.1 jmespath==0.10.0 Pygments==2.8.1 diff --git a/setup.py b/setup.py index aaf7434..7913ed0 100644 --- a/setup.py +++ b/setup.py @@ -15,9 +15,7 @@ 'jmespath', 'pygments', 'pyyaml', - 'tabulate', - # On Windows, colorama is required for legacy terminals. - 'colorama; sys_platform == "win32"' + 'tabulate' ] with open('README.rst', 'r') as f: