Init colorama only in Windows legacy terminal (#238)

Author: jiasli

knack/cli.py

@@ -11,7 +11,7 @@
 from .completion import CLICompletion
 from .output import OutputProducer
 from .log import CLILogging, get_logger
-from .util import CLIError
+from .util import CLIError, is_modern_terminal
 from .config import CLIConfig
 from .query import CLIQuery
 from .events import EVENT_CLI_PRE_EXECUTE, EVENT_CLI_SUCCESSFUL_EXECUTE, EVENT_CLI_POST_EXECUTE
@@ -21,10 +21,6 @@
 
 logger = get_logger(__name__)
 
-# Temporarily force color to be enabled even when out_file is not stdout.
-# This is only intended for testing purpose.
-_KNACK_TEST_FORCE_ENABLE_COLOR = False
-
 
 class CLI(object):  # pylint: disable=too-many-instance-attributes
     """ The main driver for the CLI """
@@ -101,6 +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 os.name == 'nt' and not is_modern_terminal()
 
     @staticmethod
     def _should_show_version(args):
@@ -207,7 +205,8 @@ 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.enable_color or _KNACK_TEST_FORCE_ENABLE_COLOR:
+            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
@@ -250,7 +249,7 @@ def invoke(self, args, initial_invocation_data=None, out_file=None):
         finally:
             self.raise_event(EVENT_CLI_POST_EXECUTE)
 
-            if self.enable_color or _KNACK_TEST_FORCE_ENABLE_COLOR:
+            if self._should_init_colorama:
                 import colorama
                 colorama.deinit()
 
@@ -274,14 +273,13 @@ def _should_enable_color(self):
             self.init_debug_log.append("Color is disabled by config.")
             return False
 
-        if 'PYCHARM_HOSTED' in os.environ:
-            if sys.stdout == sys.__stdout__ and sys.stderr == sys.__stderr__:
-                self.init_debug_log.append("Enable color in PyCharm.")
-                return True
-        else:
-            if sys.stdout.isatty() and sys.stderr.isatty() and self.out_file is sys.stdout:
-                self.init_debug_log.append("Enable color in terminal.")
-                return True
+        if sys.stdout.isatty() and sys.stderr.isatty() and self.out_file is sys.stdout:
+            self.init_debug_log.append("Enable color in terminal.")
+            return True
+
+        if 'PYCHARM_HOSTED' in os.environ and sys.stdout == sys.__stdout__ and sys.stderr == sys.__stderr__:
+            self.init_debug_log.append("Enable color in PyCharm.")
+            return True
 
         self.init_debug_log.append("Cannot enable color.")
         return False

knack/util.py

@@ -142,3 +142,24 @@ def todict(obj, post_processor=None):  # pylint: disable=too-many-return-stateme
                   if not callable(v) and not k.startswith('_')}
         return post_processor(obj, result) if post_processor else result
     return obj
+
+
+def is_modern_terminal():
+    """Detect whether the current terminal is a modern terminal that supports Unicode and
+    Console Virtual Terminal Sequences.
+
+    Currently, these terminals can be detected:
+      - VS Code terminal
+      - PyCharm
+      - Windows Terminal
+    """
+    # VS Code: https://github.com/microsoft/vscode/pull/30346
+    if os.environ.get('TERM_PROGRAM', '').lower() == 'vscode':
+        return True
+    # PyCharm: https://youtrack.jetbrains.com/issue/PY-4853
+    if 'PYCHARM_HOSTED' in os.environ:
+        return True
+    # Windows Terminal: https://github.com/microsoft/terminal/issues/1040
+    if 'WT_SESSION' in os.environ:
+        return True
+    return False

tests/test_deprecation.py

@@ -8,12 +8,11 @@
     import mock
 except ImportError:
     from unittest import mock
-from threading import Lock
 
 from knack.arguments import ArgumentsContext
-from knack.commands import CLICommand, CLICommandsLoader, CommandGroup
+from knack.commands import CLICommandsLoader, CommandGroup
 
-from tests.util import DummyCLI, redirect_io, disable_color
+from tests.util import DummyCLI, redirect_io, assert_in_multi_line, disable_color
 
 
 def example_handler(arg1, arg2=None, arg3=None):
@@ -80,7 +79,7 @@ def test_deprecate_command_group_help(self):
     cmd4 [Deprecated] : Short summary here.
 
 """.format(self.cli_ctx.name)
-        self.assertEqual(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_deprecate_command_help_hidden(self):
@@ -100,7 +99,7 @@ def test_deprecate_command_help_hidden(self):
     --arg -a           : Allowed values: 1, 2, 3.
     --arg3
 """.format(self.cli_ctx.name)
-        self.assertIn(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_deprecate_command_plain_execute(self):
@@ -211,7 +210,7 @@ def test_deprecate_command_group_help_plain(self):
     cmd1 : Short summary here.
 
 """.format(self.cli_ctx.name)
-        self.assertEqual(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_deprecate_command_group_help_hidden(self):
@@ -229,7 +228,7 @@ def test_deprecate_command_group_help_hidden(self):
     cmd1 : Short summary here.
 
 """.format(self.cli_ctx.name)
-        self.assertIn(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_deprecate_command_group_help_expiring(self):
@@ -243,7 +242,7 @@ def test_deprecate_command_group_help_expiring(self):
         This command group has been deprecated and will be removed in version '1.0.0'. Use
         'alt-group4' instead.
 """.format(self.cli_ctx.name)
-        self.assertIn(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     @disable_color
@@ -282,7 +281,7 @@ def test_deprecate_command_implicitly(self):
         This command is implicitly deprecated because command group 'group1' is deprecated and
         will be removed in a future release. Use 'alt-group1' instead.
 """.format(self.cli_ctx.name)
-        self.assertIn(expected, actual)
+        assert_in_multi_line(expected, actual)
 
 
 class TestArgumentDeprecation(unittest.TestCase):
@@ -352,7 +351,7 @@ def test_deprecate_arguments_command_help(self):
     --opt4
     --opt5
 """.format(self.cli_ctx.name)
-        self.assertTrue(actual.startswith(expected))
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_deprecate_arguments_execute(self):

tests/test_experimental.py

@@ -15,7 +15,7 @@
 from knack.arguments import ArgumentsContext
 from knack.commands import CLICommandsLoader, CommandGroup
 
-from tests.util import DummyCLI, redirect_io, remove_space
+from tests.util import DummyCLI, redirect_io, assert_in_multi_line
 
 
 def example_handler(arg1, arg2=None, arg3=None):
@@ -64,7 +64,7 @@ def test_experimental_command_implicitly_execute(self):
         self.cli_ctx.invoke('grp1 cmd1 -b b'.split())
         actual = self.io.getvalue()
         expected = "Command group 'grp1' is experimental and under development."
-        self.assertIn(remove_space(expected), remove_space(actual))
+        self.assertIn(expected, actual)
 
     @redirect_io
     def test_experimental_command_group_help(self):
@@ -83,15 +83,15 @@ def test_experimental_command_group_help(self):
     cmd1 [Experimental] : Short summary here.
 
 """.format(self.cli_ctx.name)
-        self.assertEqual(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_experimental_command_plain_execute(self):
         """ Ensure general warning displayed when running experimental command. """
         self.cli_ctx.invoke('cmd1 -b b'.split())
         actual = self.io.getvalue()
         expected = "This command is experimental and under development."
-        self.assertIn(remove_space(expected), remove_space(actual))
+        self.assertIn(expected, actual)
 
 
 class TestCommandGroupExperimental(unittest.TestCase):
@@ -136,7 +136,7 @@ def test_experimental_command_group_help_plain(self):
     cmd1 : Short summary here.
 
 """.format(self.cli_ctx.name)
-        self.assertIn(remove_space(expected), remove_space(actual))
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_experimental_command_implicitly(self):
@@ -150,7 +150,7 @@ def test_experimental_command_implicitly(self):
         Long summary here. Still long summary.
         Command group 'group1' is experimental and under development.
 """.format(self.cli_ctx.name)
-        self.assertIn(remove_space(expected), remove_space(actual))
+        assert_in_multi_line(expected, actual)
 
 
 class TestArgumentExperimental(unittest.TestCase):
@@ -193,7 +193,7 @@ def test_experimental_arguments_command_help(self):
     --arg1 [Experimental] [Required] : Arg1.
         Argument '--arg1' is experimental and under development.
 """.format(self.cli_ctx.name)
-        self.assertIn(remove_space(expected), remove_space(actual))
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_experimental_arguments_execute(self):

tests/test_preview.py

@@ -15,7 +15,7 @@
 from knack.arguments import ArgumentsContext
 from knack.commands import CLICommandsLoader, CommandGroup
 
-from tests.util import DummyCLI, redirect_io, disable_color
+from tests.util import DummyCLI, redirect_io, assert_in_multi_line, disable_color
 
 
 def example_handler(arg1, arg2=None, arg3=None):
@@ -75,7 +75,7 @@ def test_preview_command_group_help(self):
     cmd1 [Preview] : Short summary here.
 
 """.format(self.cli_ctx.name)
-        self.assertEqual(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     def test_preview_command_plain_execute(self):
@@ -170,7 +170,7 @@ def test_preview_command_group_help_plain(self):
     cmd1 : Short summary here.
 
 """.format(self.cli_ctx.name)
-        self.assertEqual(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     @disable_color
@@ -202,7 +202,7 @@ def test_preview_command_implicitly(self):
         Command group 'group1' is in preview. It may be changed/removed in a future
         release.
 """.format(self.cli_ctx.name)
-        self.assertIn(expected, actual)
+        assert_in_multi_line(expected, actual)
 
 
 class TestArgumentPreview(unittest.TestCase):
@@ -245,7 +245,7 @@ def test_preview_arguments_command_help(self):
     --arg1 [Preview] [Required] : Arg1.
         Argument '--arg1' is in preview. It may be changed/removed in a future release.
 """.format(self.cli_ctx.name)
-        self.assertIn(expected, actual)
+        assert_in_multi_line(expected, actual)
 
     @redirect_io
     @disable_color

tests/test_util.py

@@ -3,11 +3,12 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-from collections import namedtuple
 import unittest
+from collections import namedtuple
 from datetime import date, time, datetime
+from unittest import mock
 
-from knack.util import todict, to_snake_case
+from knack.util import todict, to_snake_case, is_modern_terminal
 
 
 class TestUtils(unittest.TestCase):
@@ -87,6 +88,16 @@ def test_to_snake_case_already_snake(self):
         actual = to_snake_case(the_input)
         self.assertEqual(expected, actual)
 
+    def test_is_modern_terminal(self):
+        with mock.patch.dict("os.environ", clear=True):
+            self.assertEqual(is_modern_terminal(), False)
+        with mock.patch.dict("os.environ", TERM_PROGRAM='vscode'):
+            self.assertEqual(is_modern_terminal(), True)
+        with mock.patch.dict("os.environ", PYCHARM_HOSTED='1'):
+            self.assertEqual(is_modern_terminal(), True)
+        with mock.patch.dict("os.environ", WT_SESSION='c25cb945-246a-49e5-b37a-1e4b6671b916'):
+            self.assertEqual(is_modern_terminal(), True)
+
 
 if __name__ == '__main__':
     unittest.main()

tests/util.py

@@ -7,15 +7,16 @@
     import mock
 except ImportError:
     from unittest import mock
+import logging
+import os
+import re
+import shutil
 import sys
 import tempfile
-import shutil
-import os
 from io import StringIO
-import logging
-from knack.log import CLI_LOGGER_NAME
 
 from knack.cli import CLI, CLICommandsLoader, CommandInvoker
+from knack.log import CLI_LOGGER_NAME
 
 TEMP_FOLDER_NAME = "knack_temp"
 
@@ -33,8 +34,7 @@ def wrapper(self):
         cli_logger.handlers.clear()
 
         sys.stdout = sys.stderr = self.io = StringIO()
-        with mock.patch("knack.cli._KNACK_TEST_FORCE_ENABLE_COLOR", True):
-            func(self)
+        func(self)
         self.io.close()
         sys.stdout = original_stdout
         sys.stderr = original_stderr
@@ -54,8 +54,17 @@ def wrapper(self):
     return wrapper
 
 
-def remove_space(str):
-    return str.replace(' ', '').replace('\n', '')
+def _remove_control_sequence(string):
+    return re.sub(r'\x1b[^m]+m', '', string)
+
+
+def _remove_whitespace(string):
+    return re.sub(r'\s', '', string)
+
+
+def assert_in_multi_line(sub_string, string):
+    # assert sub_string is in string, with all whitespaces, line breaks and control sequences ignored
+    assert _remove_whitespace(sub_string) in _remove_control_sequence(_remove_whitespace(string))
 
 
 class MockContext(CLI):
@@ -77,7 +86,7 @@ def get_cli_version(self):
     def __init__(self, **kwargs):
         kwargs['config_dir'] = new_temp_folder()
         super().__init__(**kwargs)
-        # Force colorama to initialize
+        # Force to enable color
         self.enable_color = True