diff --git a/knack/cli.py b/knack/cli.py index 439ed8f..337ca22 100644 --- a/knack/cli.py +++ b/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 diff --git a/knack/output.py b/knack/output.py index 13673c6..8db442a 100644 --- a/knack/output.py +++ b/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] diff --git a/tests/test_cli_scenarios.py b/tests/test_cli_scenarios.py index 63ec93f..91270d4 100644 --- a/tests/test_cli_scenarios.py +++ b/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() diff --git a/tests/test_help.py b/tests/test_help.py index e05c6b0..8d901fd 100644 --- a/tests/test_help.py +++ b/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) diff --git a/tests/test_output.py b/tests/test_output.py index 50dec07..20f2b10 100644 --- a/tests/test_output.py +++ b/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') diff --git a/tests/util.py b/tests/util.py index 8768ef4..00037e8 100644 --- a/tests/util.py +++ b/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():