From d75391a896f1efa05d066a1e0e5ee0626bf6f645 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Fri, 19 Jun 2020 16:48:47 +0800 Subject: [PATCH 1/4] Detect isatty --- knack/cli.py | 51 +++++++++++++++++++++++++++++++------ knack/output.py | 5 ++-- knack/util.py | 15 +++++++++++ tests/test_cli_scenarios.py | 33 +++++++++++++++++++++++- tests/test_help.py | 51 ++++++++++++------------------------- tests/test_output.py | 7 +++-- tests/util.py | 19 +++++++++----- 7 files changed, 124 insertions(+), 57 deletions(-) diff --git a/knack/cli.py b/knack/cli.py index 439ed8f..a63d4f6 100644 --- a/knack/cli.py +++ b/knack/cli.py @@ -12,7 +12,7 @@ from .completion import CLICompletion from .output import OutputProducer from .log import CLILogging, get_logger -from .util import CLIError +from .util import CLIError, isatty from .config import CLIConfig from .query import CLIQuery from .events import EVENT_CLI_PRE_EXECUTE, EVENT_CLI_POST_EXECUTE @@ -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,29 @@ 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) + + # 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 + # 3. stderr is a tty + # Otherwise, the output in stderr won't have LEVEL tag + # 4. out_file is stdout + conditions = (not self.config.getboolean('core', 'no_color', fallback=False), + isatty(sys.stdout), isatty(sys.stderr), self.out_file is sys.stdout) + self.enable_color = all(conditions) + # Delay showing the debug message, as logging is not initialized yet + self.init_debug_log.append("enable_color({}) = enable_color_config({}) && " + "stdout.isatty({}) && stderr.isatty({}) && out_file_is_stdout({})" + .format(self.enable_color, *conditions)) @staticmethod def _should_show_version(args): @@ -171,6 +196,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 +223,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 +266,8 @@ 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 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/knack/util.py b/knack/util.py index ae78ab1..47c8414 100644 --- a/knack/util.py +++ b/knack/util.py @@ -151,3 +151,18 @@ 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 isatty(stream): + # Code copied from + # https://github.com/tartley/colorama/blob/3d8d48a95de10be25b161c914f274ec6c41d3129/colorama/ansitowin32.py#L43-L53 + import sys + if 'PYCHARM_HOSTED' in os.environ: + if stream is not None and (stream is sys.__stdout__ or stream is sys.__stderr__): + return True + try: + stream_isatty = stream.isatty + except AttributeError: + return False + else: + return stream_isatty() diff --git a/tests/test_cli_scenarios.py b/tests/test_cli_scenarios.py index 63ec93f..3339e68 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,37 @@ def load_command_table(self, args): self.assertEqual(expected_output, mock_stdout.getvalue()) self.assertEqual(0, exit_code) + @mock.patch('sys.stderr.isatty', return_value=True) + @mock.patch('sys.stdout.isatty', return_value=True) + @mock.patch.dict('os.environ') + def test_enable_color(self, *_): + # Make sure we mock a normal terminal, instead of PyCharm terminal + os.environ.pop('PYCHARM_HOSTED', None) + cli = CLI() + self.assertEqual(cli.enable_color, True) + with mock.patch.dict("os.environ", {"CLI_CORE_NO_COLOR": "true"}): + cli = CLI() + self.assertEqual(cli.enable_color, False) + with mock.patch("sys.stderr.isatty", return_value=False): + cli = CLI() + self.assertEqual(cli.enable_color, False) + with mock.patch("sys.stdout.isatty", return_value=False): + cli = CLI() + self.assertEqual(cli.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(): From 69c24dc8370f5b31015e305520a368e2f4f6dda4 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Tue, 7 Jul 2020 14:24:37 +0800 Subject: [PATCH 2/4] Make print_init_log private --- knack/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/knack/cli.py b/knack/cli.py index a63d4f6..bdca167 100644 --- a/knack/cli.py +++ b/knack/cli.py @@ -196,7 +196,7 @@ def exception_handler(self, ex): # pylint: disable=no-self-use logger.exception(ex) return 1 - def print_init_log(self): + 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)) @@ -234,7 +234,7 @@ def invoke(self, args, initial_invocation_data=None, out_file=None): self.logging.configure(args) logger.debug('Command arguments: %s', args) - self.print_init_log() + self._print_init_log() self.raise_event(EVENT_CLI_PRE_EXECUTE) if CLI._should_show_version(args): From 6ac1274ee5dc722382a62f40720f7303681860a0 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Wed, 2 Dec 2020 16:25:42 +0800 Subject: [PATCH 3/4] Remove util.isatty --- knack/cli.py | 50 ++++++++++++++++++++++++++++++++------------------ knack/util.py | 15 --------------- 2 files changed, 32 insertions(+), 33 deletions(-) diff --git a/knack/cli.py b/knack/cli.py index bdca167..337ca22 100644 --- a/knack/cli.py +++ b/knack/cli.py @@ -12,7 +12,7 @@ from .completion import CLICompletion from .output import OutputProducer from .log import CLILogging, get_logger -from .util import CLIError, isatty +from .util import CLIError from .config import CLIConfig from .query import CLIQuery from .events import EVENT_CLI_PRE_EXECUTE, EVENT_CLI_POST_EXECUTE @@ -101,23 +101,7 @@ def __init__(self, self.init_info_log = [] self.only_show_errors = self.config.getboolean('core', 'only_show_errors', fallback=False) - - # 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 - # 3. stderr is a tty - # Otherwise, the output in stderr won't have LEVEL tag - # 4. out_file is stdout - conditions = (not self.config.getboolean('core', 'no_color', fallback=False), - isatty(sys.stdout), isatty(sys.stderr), self.out_file is sys.stdout) - self.enable_color = all(conditions) - # Delay showing the debug message, as logging is not initialized yet - self.init_debug_log.append("enable_color({}) = enable_color_config({}) && " - "stdout.isatty({}) && stderr.isatty({}) && out_file_is_stdout({})" - .format(self.enable_color, *conditions)) + self.enable_color = self._should_enable_color() @staticmethod def _should_show_version(args): @@ -271,3 +255,33 @@ def invoke(self, args, initial_invocation_data=None, out_file=None): 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/util.py b/knack/util.py index 47c8414..ae78ab1 100644 --- a/knack/util.py +++ b/knack/util.py @@ -151,18 +151,3 @@ 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 isatty(stream): - # Code copied from - # https://github.com/tartley/colorama/blob/3d8d48a95de10be25b161c914f274ec6c41d3129/colorama/ansitowin32.py#L43-L53 - import sys - if 'PYCHARM_HOSTED' in os.environ: - if stream is not None and (stream is sys.__stdout__ or stream is sys.__stderr__): - return True - try: - stream_isatty = stream.isatty - except AttributeError: - return False - else: - return stream_isatty() From db41a9f3d6caa40d6cd282d520ed67078d478dcf Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Tue, 8 Dec 2020 14:57:35 +0800 Subject: [PATCH 4/4] Update tests --- tests/test_cli_scenarios.py | 36 +++++++++++++++++++++++------------- 1 file changed, 23 insertions(+), 13 deletions(-) diff --git a/tests/test_cli_scenarios.py b/tests/test_cli_scenarios.py index 3339e68..91270d4 100644 --- a/tests/test_cli_scenarios.py +++ b/tests/test_cli_scenarios.py @@ -126,23 +126,33 @@ def load_command_table(self, args): self.assertEqual(expected_output, mock_stdout.getvalue()) self.assertEqual(0, exit_code) - @mock.patch('sys.stderr.isatty', return_value=True) - @mock.patch('sys.stdout.isatty', return_value=True) + @mock.patch('sys.stderr.isatty') + @mock.patch('sys.stdout.isatty') @mock.patch.dict('os.environ') - def test_enable_color(self, *_): + 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() - self.assertEqual(cli.enable_color, True) - with mock.patch.dict("os.environ", {"CLI_CORE_NO_COLOR": "true"}): - cli = CLI() - self.assertEqual(cli.enable_color, False) - with mock.patch("sys.stderr.isatty", return_value=False): - cli = CLI() - self.assertEqual(cli.enable_color, False) - with mock.patch("sys.stdout.isatty", return_value=False): - cli = CLI() - self.assertEqual(cli.enable_color, False) + + # 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):