[Color] Detect isatty when enabling color

knack/cli.py

@@ -22,6 +22,10 @@
 
 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 """
@@ -91,8 +95,13 @@ def __init__(self,
         self.output = self.output_cls(cli_ctx=self)
         self.result = None
         self.query = query_cls(cli_ctx=self)
+
+        # As logging is initialized in `invoke`, call `logger.debug` or `logger.info` here won't work.
+        self.init_debug_log = []
+        self.init_info_log = []
+
         self.only_show_errors = self.config.getboolean('core', 'only_show_errors', fallback=False)
-        self.enable_color = not self.config.getboolean('core', 'no_color', fallback=False)
+        self.enable_color = self._should_enable_color()
 
     @staticmethod
     def _should_show_version(args):
@@ -171,6 +180,15 @@ def exception_handler(self, ex):  # pylint: disable=no-self-use
             logger.exception(ex)
         return 1
 
+    def _print_init_log(self):
+        """Print the debug/info log from CLI.__init__"""
+        if self.init_debug_log:
+            logger.debug('__init__ debug log:\n%s', '\n'.join(self.init_debug_log))
+            self.init_debug_log.clear()
+        if self.init_info_log:
+            logger.info('__init__ info log:\n%s', '\n'.join(self.init_info_log))
+            self.init_info_log.clear()
+
     def invoke(self, args, initial_invocation_data=None, out_file=None):
         """ Invoke a command.
 
@@ -189,18 +207,18 @@ def invoke(self, args, initial_invocation_data=None, out_file=None):
             raise TypeError('args should be a list or tuple.')
         exit_code = 0
         try:
-            if self.enable_color:
+            out_file = out_file or self.out_file
+            if out_file is sys.stdout and self.enable_color or _KNACK_TEST_FORCE_ENABLE_COLOR:
                 import colorama
                 colorama.init()
-                if self.out_file == sys.__stdout__:
-                    # point out_file to the new sys.stdout which is overwritten by colorama
-                    self.out_file = sys.stdout
+                # point out_file to the new sys.stdout which is overwritten by colorama
+                out_file = sys.stdout
 
             args = self.completion.get_completion_args() or args
-            out_file = out_file or self.out_file
 
             self.logging.configure(args)
             logger.debug('Command arguments: %s', args)
+            self._print_init_log()
 
             self.raise_event(EVENT_CLI_PRE_EXECUTE)
             if CLI._should_show_version(args):
@@ -232,7 +250,38 @@ def invoke(self, args, initial_invocation_data=None, out_file=None):
         finally:
             self.raise_event(EVENT_CLI_POST_EXECUTE)
 
-            if self.enable_color:
+            if self.enable_color or _KNACK_TEST_FORCE_ENABLE_COLOR:
+                import colorama
                 colorama.deinit()
 
         return exit_code
+
+    def _should_enable_color(self):
+        # When run in a normal terminal, color is only enabled when all conditions are met:
+        #   1. [core] no_color config is not set
+        #   2. stdout is a tty
+        #      - 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
+
+        no_color_config = self.config.getboolean('core', 'no_color', fallback=False)
+        # If color is disabled by config explicitly, never enable color
+        if no_color_config:
+            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
+
+        self.init_debug_log.append("Cannot enable color.")
+        return False

knack/output.py

@@ -8,7 +8,6 @@
 import errno
 import json
 import traceback
-import sys
 from collections import OrderedDict
 from six import StringIO, text_type, u, string_types
 
@@ -162,9 +161,9 @@ def out(self, obj, formatter=None, out_file=None):  # pylint: disable=no-self-us
 
     def get_formatter(self, format_type):  # pylint: disable=no-self-use
         # remove color if stdout is not a tty
-        if not sys.stdout.isatty() and format_type == 'jsonc':
+        if not self.cli_ctx.enable_color and format_type == 'jsonc':
             return OutputProducer._FORMAT_DICT['json']
-        if not sys.stdout.isatty() and format_type == 'yamlc':
+        if not self.cli_ctx.enable_color and format_type == 'yamlc':
             return OutputProducer._FORMAT_DICT['yaml']
         return OutputProducer._FORMAT_DICT[format_type]
 

tests/test_cli_scenarios.py

@@ -17,7 +17,7 @@
 from knack import CLI
 from knack.commands import CLICommand, CLICommandsLoader
 from knack.invocation import CommandInvoker
-from tests.util import MockContext
+from tests.util import MockContext, redirect_io
 
 
 class TestCLIScenarios(unittest.TestCase):
@@ -126,6 +126,47 @@ def load_command_table(self, args):
         self.assertEqual(expected_output, mock_stdout.getvalue())
         self.assertEqual(0, exit_code)
 
+    @mock.patch('sys.stderr.isatty')
+    @mock.patch('sys.stdout.isatty')
+    @mock.patch.dict('os.environ')
+    def test_should_enable_color(self, stdout_isatty_mock, stderr_isatty_mock):
+        # Make sure we mock a normal terminal, instead of PyCharm terminal
+        os.environ.pop('PYCHARM_HOSTED', None)
+        cli = CLI()
+
+        # Color is turned on by default
+        stdout_isatty_mock.return_value = True
+        stderr_isatty_mock.return_value = True
+        self.assertEqual(cli._should_enable_color(), True)
+
+        # Color is turned off with a main switch
+        os.environ['CLI_CORE_NO_COLOR'] = 'yes'
+        self.assertEqual(cli._should_enable_color(), False)
+        del os.environ['CLI_CORE_NO_COLOR']
+
+        # Mock stderr is not a TTY
+        stdout_isatty_mock.return_value = False
+        stderr_isatty_mock.return_value = True
+        self.assertEqual(cli._should_enable_color(), False)
+
+        # Mock stdout is not a TTY
+        stdout_isatty_mock.return_value = True
+        stderr_isatty_mock.return_value = False
+        self.assertEqual(cli._should_enable_color(), False)
+
+    @redirect_io
+    def test_init_log(self):
+        class MyCLI(CLI):
+            def __init__(self, **kwargs):
+                super(MyCLI, self).__init__(**kwargs)
+                self.init_debug_log.append("init debug log: 6aa19a11")
+                self.init_info_log.append("init info log: b0746f58")
+        cli = MyCLI()
+        cli.invoke(["--debug"])
+        actual = self.io.getvalue()
+        self.assertIn("6aa19a11", actual)
+        self.assertIn("b0746f58", actual)
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_help.py

@@ -5,32 +5,13 @@
 
 import sys
 import unittest
-import mock
-from six import StringIO
-
 
-from knack.help import ArgumentGroupRegistry, HelpObject
+import mock
 from knack.arguments import ArgumentsContext
-from knack.commands import CLICommand, CLICommandsLoader, CommandGroup
+from knack.commands import CLICommandsLoader, CommandGroup
 from knack.events import EVENT_PARSER_GLOBAL_CREATE
-from knack.invocation import CommandInvoker
-
-from tests.util import MockContext, DummyCLI
-
-io = {}
-
-
-def redirect_io(func):
-    def wrapper(self):
-        global io
-        old_stdout = sys.stdout
-        old_stderr = sys.stderr
-        sys.stdout = sys.stderr = io = StringIO()
-        func(self)
-        io.close()
-        sys.stdout = old_stdout
-        sys.stderr = old_stderr
-    return wrapper
+from knack.help import ArgumentGroupRegistry, HelpObject
+from tests.util import DummyCLI, redirect_io
 
 
 def example_handler(arg1, arg2=None, arg3=None):
@@ -206,7 +187,7 @@ def test_choice_list_with_ints(self):
         """ Ensure choice_list works with integer lists. """
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n1 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = 'Allowed values: 1, 2, 3'
         self.assertIn(expected, actual)
 
@@ -227,7 +208,7 @@ def test_help_long_and_short_description_from_docstring(self):
 
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n1 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = '\nCommand\n    {} n1 : Short summary here.\n        Long summary here. Still long summary.'.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -238,7 +219,7 @@ def test_help_long_and_short_description_from_yaml(self):
 
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n2 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = '\nCommand\n    {} n2 : YAML short summary.\n        YAML long summary. More summary.'.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -248,7 +229,7 @@ def test_help_long_description_multi_line(self):
 
         with self.assertRaises(SystemExit):
             self.cli_ctx.invoke('n3 -h'.split())
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = '\nCommand\n    {} n3 : Short summary here.\n        Line1\n        line2.\n'.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -269,7 +250,7 @@ def test_help_params_documentations(self):
         Paragraph(s).
     --foobar3            : The foobar3.
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -307,7 +288,7 @@ def test_help_full_documentations(self):
         example details
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertTrue(actual.startswith(expected))
 
@@ -338,7 +319,7 @@ def test_help_with_param_specified(self):
     --verbose          : Increase logging verbosity. Use --debug for full debug logs.
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 
@@ -358,7 +339,7 @@ def test_help_group_children(self):
     beta
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 
@@ -377,7 +358,7 @@ def test_help_missing_params(self):
             with self.assertRaises(SystemExit):
                 self.cli_ctx.invoke('n1 -a 1 --arg 2'.split())
 
-            actual = io.getvalue()
+            actual = self.io.getvalue()
             self.assertTrue('required' in actual and '-b' in actual)
 
     @redirect_io
@@ -395,7 +376,7 @@ def test_help_extra_params(self):
             with self.assertRaises(SystemExit):
                 self.cli_ctx.invoke('n1 -a 1 -b c -c extra'.split())
 
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = 'unrecognized arguments: -c extra'
         self.assertIn(expected, actual)
 
@@ -415,7 +396,7 @@ def test_help_group_help(self):
     n1 : This module does xyz one-line or so.
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = expected.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 
@@ -455,7 +436,7 @@ def register_globals(_, **kwargs):
     --verbose          : Increase logging verbosity. Use --debug for full debug logs.
 
 """
-        actual = io.getvalue()
+        actual = self.io.getvalue()
         expected = s.format(self.cli_ctx.name)
         self.assertEqual(actual, expected)
 

tests/test_output.py

@@ -277,17 +277,16 @@ def test_output_format_ordereddict_list_not_sorted(self):
         result = format_tsv(CommandResultItem([obj1, obj2]))
         self.assertEqual(result, '1\t2\n3\t4\n')
 
-    @mock.patch('sys.stdout.isatty', autospec=True)
-    def test_remove_color_no_tty(self, mock_isatty):
+    def test_remove_color_no_tty(self):
         output_producer = OutputProducer(cli_ctx=self.mock_ctx)
 
-        mock_isatty.return_value = False
+        self.mock_ctx.enable_color = False
         formatter = output_producer.get_formatter('jsonc')
         self.assertEqual(formatter, format_json)
         formatter = output_producer.get_formatter('yamlc')
         self.assertEqual(formatter, format_yaml)
 
-        mock_isatty.return_value = True
+        self.mock_ctx.enable_color = True
         formatter = output_producer.get_formatter('jsonc')
         self.assertEqual(formatter, format_json_color)
         formatter = output_producer.get_formatter('yamlc')

tests/util.py

@@ -26,18 +26,23 @@ def redirect_io(func):
     original_stdout = sys.stdout
 
     def wrapper(self):
+        # Ensure a clean startup - no log handlers
+        root_logger = logging.getLogger()
+        cli_logger = logging.getLogger(CLI_LOGGER_NAME)
+        root_logger.handlers.clear()
+        cli_logger.handlers.clear()
+
         sys.stdout = sys.stderr = self.io = StringIO()
-        func(self)
+        with mock.patch("knack.cli._KNACK_TEST_FORCE_ENABLE_COLOR", True):
+            func(self)
         self.io.close()
-        sys.stdout = original_stderr
+        sys.stdout = original_stdout
         sys.stderr = original_stderr
 
         # Remove the handlers added by CLI, so that the next invoke call init them again with the new stderr
         # Otherwise, the handlers will write to a closed StringIO from a preview test
-        root_logger = logging.getLogger()
-        cli_logger = logging.getLogger(CLI_LOGGER_NAME)
-        root_logger.handlers = root_logger.handlers[:-1]
-        cli_logger.handlers = cli_logger.handlers[:-1]
+        root_logger.handlers.clear()
+        cli_logger.handlers.clear()
     return wrapper
 
 
@@ -72,6 +77,8 @@ def get_cli_version(self):
     def __init__(self, **kwargs):
         kwargs['config_dir'] = new_temp_folder()
         super(DummyCLI, self).__init__(**kwargs)
+        # Force colorama to initialize
+        self.enable_color = True
 
 
 def new_temp_folder():