Enable Virtual Terminal mode on legacy Windows terminal to support ANSI escape sequences

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

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

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

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: