diff --git a/knack/output.py b/knack/output.py index e5e07ba..110f1e0 100644 --- a/knack/output.py +++ b/knack/output.py @@ -8,8 +8,10 @@ import errno import json import traceback +import sys from collections import OrderedDict from six import StringIO, text_type, u, string_types +import yaml from .util import CLIError, CommandResultItem, CtxTypeError from .events import EVENT_INVOKER_POST_PARSE_ARGS, EVENT_PARSER_GLOBAL_CREATE @@ -45,6 +47,23 @@ def format_json_color(obj): return highlight(format_json(obj), lexers.JsonLexer(), formatters.TerminalFormatter()) # pylint: disable=no-member +def format_yaml(obj): + try: + return yaml.safe_dump(obj.result, default_flow_style=False, allow_unicode=True) + except yaml.representer.RepresenterError: + # yaml.safe_dump fails when obj.result is an OrderedDict. knack's --query implementation converts the result to an OrderedDict. https://github.com/microsoft/knack/blob/af674bfea793ff42ae31a381a21478bae4b71d7f/knack/query.py#L46. # pylint: disable=line-too-long + return yaml.safe_dump(json.loads(json.dumps(obj.result)), default_flow_style=False, allow_unicode=True) + + +def format_yaml_color(obj): + from pygments import highlight, lexers, formatters + return highlight(format_yaml(obj), lexers.YamlLexer(), formatters.TerminalFormatter()) # pylint: disable=no-member + + +def format_none(_): + return "" + + def format_table(obj): result = obj.result try: @@ -78,8 +97,11 @@ class OutputProducer(object): _FORMAT_DICT = { 'json': format_json, 'jsonc': format_json_color, + 'yaml': format_yaml, + 'yamlc': format_yaml_color, 'table': format_table, 'tsv': format_tsv, + 'none': format_none, } @staticmethod @@ -142,6 +164,11 @@ def out(self, obj, formatter=None, out_file=None): # pylint: disable=no-self-us file=out_file, end='') 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': + return OutputProducer._FORMAT_DICT['json'] + if not sys.stdout.isatty() and format_type == 'yamlc': + return OutputProducer._FORMAT_DICT['yaml'] return OutputProducer._FORMAT_DICT[format_type] diff --git a/tests/test_help.py b/tests/test_help.py index e5c09b9..e1ca100 100644 --- a/tests/test_help.py +++ b/tests/test_help.py @@ -295,7 +295,8 @@ def test_help_full_documentations(self): Global Arguments --debug : Increase logging verbosity to show all debug logs. --help -h : Show this help message and exit. - --output -o : Output format. Allowed values: json, jsonc, table, tsv. Default: json. + --output -o : Output format. Allowed values: json, jsonc, none, table, tsv, yaml, + yamlc. Default: json. --query : JMESPath query string. See http://jmespath.org/ for more information and examples. --verbose : Increase logging verbosity. Use --debug for full debug logs. @@ -328,7 +329,8 @@ def test_help_with_param_specified(self): Global Arguments --debug : Increase logging verbosity to show all debug logs. --help -h : Show this help message and exit. - --output -o : Output format. Allowed values: json, jsonc, table, tsv. Default: json. + --output -o : Output format. Allowed values: json, jsonc, none, table, tsv, yaml, yamlc. + Default: json. --query : JMESPath query string. See http://jmespath.org/ for more information and examples. --verbose : Increase logging verbosity. Use --debug for full debug logs. @@ -443,7 +445,8 @@ def register_globals(_, **kwargs): --debug : Increase logging verbosity to show all debug logs. --exampl : This is a new global argument. --help -h : Show this help message and exit. - --output -o : Output format. Allowed values: json, jsonc, table, tsv. Default: json. + --output -o : Output format. Allowed values: json, jsonc, none, table, tsv, yaml, yamlc. + Default: json. --query : JMESPath query string. See http://jmespath.org/ for more information and examples. --verbose : Increase logging verbosity. Use --debug for full debug logs. diff --git a/tests/test_output.py b/tests/test_output.py index d2d7a4a..50dec07 100644 --- a/tests/test_output.py +++ b/tests/test_output.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- +# -*- coding: utf-8 -*- # -------------------------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. @@ -9,10 +9,12 @@ from __future__ import print_function import unittest +import mock from collections import OrderedDict from six import StringIO -from knack.output import OutputProducer, format_json, format_table, format_tsv +from knack.output import OutputProducer, format_json, format_json_color, format_yaml, format_yaml_color, \ + format_table, format_tsv from knack.util import CommandResultItem, normalize_newlines from tests.util import MockContext @@ -30,6 +32,8 @@ def test_cli_ctx_type_error(self): with self.assertRaises(TypeError): OutputProducer(cli_ctx=object()) + # JSON output tests + def test_out_json_valid(self): """ The JSON output when the input is a dict should be the dict serialized to JSON @@ -89,6 +93,60 @@ def test_out_json_non_ASCII(self): "active": true, "contents": "生活很糟糕" } +""")) + + # YAML output tests + + def test_out_yaml_valid(self): + """ + Test Dict serialized to YAML + """ + output_producer = OutputProducer(cli_ctx=self.mock_ctx) + output_producer.out(CommandResultItem({'active': True, 'id': '0b1f6472'}), + formatter=format_yaml, out_file=self.io) + self.assertEqual(normalize_newlines(self.io.getvalue()), normalize_newlines( + """active: true +id: 0b1f6472 +""")) + + def test_out_yaml_from_ordered_dict(self): + """ + Test OrderedDict serialized to YAML + """ + output_producer = OutputProducer(cli_ctx=self.mock_ctx) + output_producer.out(CommandResultItem(OrderedDict({'active': True, 'id': '0b1f6472'})), + formatter=format_yaml, out_file=self.io) + self.assertEqual(normalize_newlines(self.io.getvalue()), normalize_newlines( + """active: true +id: 0b1f6472 +""")) + + def test_out_yaml_byte(self): + output_producer = OutputProducer(cli_ctx=self.mock_ctx) + output_producer.out(CommandResultItem({'active': True, 'contents': b'0b1f6472'}), + formatter=format_yaml, out_file=self.io) + self.assertEqual(normalize_newlines(self.io.getvalue()), normalize_newlines( + """active: true +contents: !!binary | + MGIxZjY0NzI= +""")) + + def test_out_yaml_byte_empty(self): + output_producer = OutputProducer(cli_ctx=self.mock_ctx) + output_producer.out(CommandResultItem({'active': True, 'contents': b''}), + formatter=format_yaml, out_file=self.io) + self.assertEqual(normalize_newlines(self.io.getvalue()), normalize_newlines( + """active: true +contents: !!binary "" +""")) + + def test_out_yaml_non_ASCII(self): + output_producer = OutputProducer(cli_ctx=self.mock_ctx) + output_producer.out(CommandResultItem({'active': True, 'contents': 'こんにちは'}), + formatter=format_yaml, out_file=self.io) + self.assertEqual(normalize_newlines(self.io.getvalue()), normalize_newlines( + """active: true +contents: こんにちは """)) # TABLE output tests @@ -219,6 +277,22 @@ 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): + output_producer = OutputProducer(cli_ctx=self.mock_ctx) + + mock_isatty.return_value = 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 + formatter = output_producer.get_formatter('jsonc') + self.assertEqual(formatter, format_json_color) + formatter = output_producer.get_formatter('yamlc') + self.assertEqual(formatter, format_yaml_color) + if __name__ == '__main__': unittest.main()