bpo-29636: Add --(no-)indent arguments to json.tool (GH-345)

2 files changed, 64 insertions, 7 deletions

diff --git a/Lib/json/tool.py b/Lib/json/tool.py
index b3ef992..6c687d7 100644
--- a/Lib/json/tool.py
+++ b/Lib/json/tool.py
@@ -30,20 +30,36 @@ def main():
                         help='sort the output of dictionaries alphabetically by key')
     parser.add_argument('--json-lines', action='store_true', default=False,
                         help='parse input using the jsonlines format')
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument('--indent', default=4, type=int,
+                       help='separate items with newlines and use this number '
+                       'of spaces for indentation')
+    group.add_argument('--tab', action='store_const', dest='indent',
+                       const='\t', help='separate items with newlines and use '
+                       'tabs for indentation')
+    group.add_argument('--no-indent', action='store_const', dest='indent',
+                       const=None,
+                       help='separate items with spaces rather than newlines')
+    group.add_argument('--compact', action='store_true',
+                       help='suppress all whitespace separation (most compact)')
     options = parser.parse_args()
 
-    infile = options.infile
-    outfile = options.outfile
-    sort_keys = options.sort_keys
-    json_lines = options.json_lines
-    with infile, outfile:
+    dump_args = {
+        'sort_keys': options.sort_keys,
+        'indent': options.indent,
+    }
+    if options.compact:
+        dump_args['indent'] = None
+        dump_args['separators'] = ',', ':'
+
+    with options.infile as infile, options.outfile as outfile:
         try:
-            if json_lines:
+            if options.json_lines:
                 objs = (json.loads(line) for line in infile)
             else:
                 objs = (json.load(infile), )
             for obj in objs:
-                json.dump(obj, outfile, sort_keys=sort_keys, indent=4)
+                json.dump(obj, outfile, **dump_args)
                 outfile.write('\n')
         except ValueError as e:
             raise SystemExit(e)
diff --git a/Lib/test/test_json/test_tool.py b/Lib/test/test_json/test_tool.py
index 1e95bc7..81d179c 100644
--- a/Lib/test/test_json/test_tool.py
+++ b/Lib/test/test_json/test_tool.py
@@ -134,3 +134,44 @@ class TestTool(unittest.TestCase):
         self.assertEqual(out.splitlines(),
                          self.expect_without_sort_keys.encode().splitlines())
         self.assertEqual(err, b'')
+
+    def test_indent(self):
+        json_stdin = b'[1, 2]'
+        expect = textwrap.dedent('''\
+        [
+          1,
+          2
+        ]
+        ''').encode()
+        args = sys.executable, '-m', 'json.tool', '--indent', '2'
+        with Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE) as proc:
+            json_stdout, err = proc.communicate(json_stdin)
+        self.assertEqual(expect.splitlines(), json_stdout.splitlines())
+        self.assertEqual(err, b'')
+
+    def test_no_indent(self):
+        json_stdin = b'[1,\n2]'
+        expect = b'[1, 2]'
+        args = sys.executable, '-m', 'json.tool', '--no-indent'
+        with Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE) as proc:
+            json_stdout, err = proc.communicate(json_stdin)
+        self.assertEqual(expect.splitlines(), json_stdout.splitlines())
+        self.assertEqual(err, b'')
+
+    def test_tab(self):
+        json_stdin = b'[1, 2]'
+        expect = b'[\n\t1,\n\t2\n]\n'
+        args = sys.executable, '-m', 'json.tool', '--tab'
+        with Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE) as proc:
+            json_stdout, err = proc.communicate(json_stdin)
+        self.assertEqual(expect.splitlines(), json_stdout.splitlines())
+        self.assertEqual(err, b'')
+
+    def test_compact(self):
+        json_stdin = b'[ 1 ,\n 2]'
+        expect = b'[1,2]'
+        args = sys.executable, '-m', 'json.tool', '--compact'
+        with Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE) as proc:
+            json_stdout, err = proc.communicate(json_stdin)
+        self.assertEqual(expect.splitlines(), json_stdout.splitlines())
+        self.assertEqual(err, b'')