Add support for json lines

python · remilapeyre · Jun 22, 2018 · Jun 26, 2018 · Jun 27, 2018 · Jun 29, 2018
commit fe960a5bb6ced5b6b3fd203ff792d00f9886d9ea
diff --git a/Lib/json/tool.py b/Lib/json/tool.py
@@ -14,43 +14,69 @@
 import json
 import sys
 
+def _read(infile, json_lines):
+            try:
+                if json_lines:
+                    return (json.loads(line) for line in infile)
+                else:
+                    return (json.load(infile), )
+            except ValueError as e:
+                raise SystemExit(e)
+
+def _open_outfile(outfile, parser):
+    try:
+        return sys.stdout if outfile == '-' else open(outfile, 'w')
+    except IOError as e:
+        parser.error(f"can't open '{outfile}': {str(e)}")
+
+def _write(parser, outfile, objs, sort_keys):
+    for obj in objs:
+        json.dump(obj, outfile, sort_keys=sort_keys, indent=4)
+        outfile.write('\n')
 
 def main():
     prog = 'python -m json.tool'
     description = ('A simple command line interface for json module '
                    'to validate and pretty-print JSON objects.')
     parser = argparse.ArgumentParser(prog=prog, description=description)
-    parser.add_argument('infile', nargs='?', type=argparse.FileType(),
+    parser.add_argument('infile', nargs='?', type=argparse.FileType(), default=sys.stdin,
                         help='a JSON file to be validated or pretty-printed')
     parser.add_argument('outfile', nargs='?', default='-',
                         help='write the output of infile to outfile')
     parser.add_argument('--sort-keys', action='store_true', default=False,
                         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')
+    parser.add_argument('-i', '--in-place', action='store_true', default=False,
+                        help='edit the file in-place')
     options = parser.parse_args()
 
-    infile = options.infile or sys.stdin
+    if options.in_place:
+        if options.outfile != '-':
+            parser.error('outfile cannot be set when -i / --in-place is used')
+        if options.infile is sys.stdin:
+            parser.error('infile must be set when -i / --in-place is used')
+        options.outfile = options.infile.name
+
+    outfile = options.outfile
+    infile = options.infile
     sort_keys = options.sort_keys
     json_lines = options.json_lines
-    with infile:
-        try:
-            if json_lines:
-                objs = tuple(json.loads(line) for line in infile)
-            else:
-                objs = (json.load(infile), )
-        except ValueError as e:
-            raise SystemExit(e)
 
-    try:
-        outfile = sys.stdout if options.outfile == '-' else open(options.outfile, 'w')
-    except IOError as e:
-        parser.error(f"can't open '{options.outfile}': {str(e)}")
-    with outfile:
-        for obj in objs:
-            json.dump(obj, outfile, sort_keys=sort_keys, indent=4)
-            outfile.write('\n')
+    if options.in_place:
+        with infile:
+            objs = tuple(_read(infile, json_lines))
+
+        outfile = _open_outfile(outfile, parser)
+
+        with outfile:
+            _write(parser, outfile, objs, sort_keys)
 
+    else:
+        outfile = _open_outfile(outfile, parser)
+        with infile, outfile:
+            objs = _read(infile, json_lines)
+            _write(parser, outfile, objs, sort_keys)
 
 if __name__ == '__main__':
     main()
diff --git a/Lib/test/test_json/test_tool.py b/Lib/test/test_json/test_tool.py
@@ -92,11 +92,13 @@ def test_stdin_stdout(self):
         self.assertEqual(out.splitlines(), self.expect.encode().splitlines())
         self.assertEqual(err, b'')
 
-    def _create_infile(self):
+    def _create_infile(self, data=None):
+        if data is None:
+            data = self.data
         infile = support.TESTFN
         with open(infile, "w") as fp:
             self.addCleanup(os.remove, infile)
-            fp.write(self.data)
+            fp.write(data)
         return infile
 
     def test_infile_stdout(self):
@@ -117,7 +119,7 @@ def test_infile_outfile(self):
 
     def test_infile_same_outfile(self):
         infile = self._create_infile()
-        rc, out, err = assert_python_ok('-m', 'json.tool', infile, infile)
+        rc, out, err = assert_python_ok('-m', 'json.tool', '-i', infile)
         self.assertEqual(out, b'')
         self.assertEqual(err, b'')
 
@@ -130,16 +132,38 @@ def test_unavailable_outfile(self):
 
     def test_jsonlines(self):
         args = sys.executable, '-m', 'json.tool', '--json-lines'
-        with Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE) as proc:
-            out, err = proc.communicate(self.jsonlines_raw.encode())
-        self.assertEqual(out.splitlines(), self.jsonlines_expect.encode().splitlines())
+        proc = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+        out, err = proc.communicate(self.jsonlines_raw.encode())
+        proc.wait()
+        self.assertEqual(proc.returncode, 0)
         self.assertEqual(err, b'')
+        self.assertEqual(out.splitlines(), self.jsonlines_expect.encode().splitlines())
 
     def test_help_flag(self):
         rc, out, err = assert_python_ok('-m', 'json.tool', '-h')
         self.assertTrue(out.startswith(b'usage: '))
         self.assertEqual(err, b'')
 
+    def test_inplace_flag(self):
+        rc, out, err = assert_python_failure('-m', 'json.tool', '-i')
+        self.assertEqual(out, b'')
+        self.assertIn(b"error: infile must be set when -i / --in-place is used", err)
+
+        rc, out, err = assert_python_failure('-m', 'json.tool', '-i', '-')
+        self.assertEqual(out, b'')
+        self.assertIn(b"error: infile must be set when -i / --in-place is used", err)
+
+        infile = self._create_infile()
+        rc, out, err = assert_python_failure('-m', 'json.tool', '-i', infile, 'test.json')
+        self.assertEqual(out, b'')
+        self.assertIn(b"error: outfile cannot be set when -i / --in-place is used", err)
+
+    def test_inplace_jsonlines(self):
+        infile = self._create_infile(data=self.jsonlines_raw)
+        rc, out, err = assert_python_ok('-m', 'json.tool', '--json-lines', '-i', infile)
+        self.assertEqual(out, b'')
+        self.assertEqual(err, b'')
+
     def test_sort_keys_flag(self):
         infile = self._create_infile()
         rc, out, err = assert_python_ok('-m', 'json.tool', '--sort-keys', infile)
@@ -156,22 +180,24 @@ def test_no_fd_leak_infile_outfile(self):
                 json.tool.main()
 
         os.unlink(infile + '.out')
-        self.assertEqual(opened, closed)
-        self.assertEqual(len(opened), len(closed))
+        self.assertEqual(set(opened), set(closed))
+        self.assertEqual(len(opened), 2)
+        self.assertEqual(len(opened), 2)
 
     def test_no_fd_leak_same_infile_outfile(self):
         infile = self._create_infile()
         closed, opened, open = mock_open()
         with mock.patch('builtins.open', side_effect=open):
-            with mock.patch.object(sys, 'argv', ['tool.py', infile, infile]):
+            with mock.patch.object(sys, 'argv', ['tool.py', '-i', infile]):
                 try:
                     import json.tool
                     json.tool.main()
                 except SystemExit:
                     pass
 
         self.assertEqual(opened, closed)
-        self.assertEqual(len(opened), len(closed))
+        self.assertEqual(len(opened), 2)
+        self.assertEqual(len(opened), 2)
 
 def mock_open():
     closed = []

diff --git a/Misc/NEWS.d/next/Library/2018-06-27-14-46-03.bpo-33927.felCrI.rst b/Misc/NEWS.d/next/Library/2018-06-27-14-46-03.bpo-33927.felCrI.rst
@@ -1,2 +1,2 @@
-`json.tool` can now take the same file as input and ouput. Patch by Rémi
-Lapeyre.
+`json.tool` can now take the same file as input and ouput with the `--in-place`
+flag. Patch by Rémi Lapeyre.