bpo-33927: Add support for same infile and outfile to json.tool

Lib/json/tool.py

@@ -14,6 +14,28 @@
 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:
+        if outfile == '-':
+            return sys.stdout
+        else:
+            return open(outfile, 'w', encoding='utf-8')
+    except IOError as e:
+        parser.error(f"can't open '{outfile}': {str(e)}")
+
+def _write(outfile, objs, **kwargs):
+    for obj in objs:
+        json.dump(obj, outfile, **kwargs)
+        outfile.write('\n')
 
 def main():
     prog = 'python -m json.tool'
@@ -24,10 +46,8 @@ def main():
                         type=argparse.FileType(encoding="utf-8"),
                         help='a JSON file to be validated or pretty-printed',
                         default=sys.stdin)
-    parser.add_argument('outfile', nargs='?',
-                        type=argparse.FileType('w', encoding="utf-8"),
-                        help='write the output of infile to outfile',
-                        default=sys.stdout)
+    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('--no-ensure-ascii', dest='ensure_ascii', action='store_false',
@@ -47,8 +67,18 @@ def main():
                        help='separate items with spaces rather than newlines')
     group.add_argument('--compact', action='store_true',
                        help='suppress all whitespace separation (most compact)')
+    parser.add_argument('-i', '--in-place', action='store_true', default=False,
+                        help='edit the file in-place')
     options = parser.parse_args()
 
+
+    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
+
     dump_args = {
         'sort_keys': options.sort_keys,
         'indent': options.indent,
@@ -58,18 +88,18 @@ def main():
         dump_args['indent'] = None
         dump_args['separators'] = ',', ':'
 
-    with options.infile as infile, options.outfile as outfile:
-        try:
-            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, **dump_args)
-                outfile.write('\n')
-        except ValueError as e:
-            raise SystemExit(e)
+    if options.in_place:
+        with options.infile as infile:
+            objs = tuple(_read(infile, options.json_lines))
+
+        with _open_outfile(options.outfile, parser) as outfile:
+            _write(outfile, objs, **dump_args)
 
+    else:
+        outfile = _open_outfile(options.outfile, parser)
+        with options.infile as infile, outfile:
+            objs = _read(infile, options.json_lines)
+            _write(outfile, objs, **dump_args)
 
 if __name__ == '__main__':
     try:

Lib/test/test_json/test_tool.py

@@ -4,9 +4,12 @@
 import textwrap
 import unittest
 import subprocess
+import io
+import types
 
 from test import support
-from test.support.script_helper import assert_python_ok
+from test.support.script_helper import assert_python_ok, assert_python_failure
+from unittest import mock
 
 
 class TestTool(unittest.TestCase):
@@ -100,7 +103,6 @@ def _create_infile(self, data=None):
     def test_infile_stdout(self):
         infile = self._create_infile()
         rc, out, err = assert_python_ok('-m', 'json.tool', infile)
-        self.assertEqual(rc, 0)
         self.assertEqual(out.splitlines(), self.expect.encode().splitlines())
         self.assertEqual(err, b'')
 
@@ -126,10 +128,22 @@ def test_infile_outfile(self):
         self.addCleanup(os.remove, outfile)
         with open(outfile, "r") as fp:
             self.assertEqual(fp.read(), self.expect)
-        self.assertEqual(rc, 0)
         self.assertEqual(out, b'')
         self.assertEqual(err, b'')
 
+    def test_infile_same_outfile(self):
+        infile = self._create_infile()
+        rc, out, err = assert_python_ok('-m', 'json.tool', '-i', infile)
+        self.assertEqual(out, b'')
+        self.assertEqual(err, b'')
+
+    def test_unavailable_outfile(self):
+        infile = self._create_infile()
+        rc, out, err = assert_python_failure('-m', 'json.tool', infile, '/bla/outfile')
+        self.assertEqual(rc, 2)
+        self.assertEqual(out, b'')
+        self.assertIn(b"error: can't open '/bla/outfile': [Errno 2]", err)
+
     def test_jsonlines(self):
         args = sys.executable, '-m', 'json.tool', '--json-lines'
         process = subprocess.run(args, input=self.jsonlines_raw, capture_output=True, text=True, check=True)
@@ -138,18 +152,64 @@ def test_jsonlines(self):
 
     def test_help_flag(self):
         rc, out, err = assert_python_ok('-m', 'json.tool', '-h')
-        self.assertEqual(rc, 0)
         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)
-        self.assertEqual(rc, 0)
         self.assertEqual(out.splitlines(),
                          self.expect_without_sort_keys.encode().splitlines())
         self.assertEqual(err, b'')
 
+    def test_no_fd_leak_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 + '.out']):
+                import json.tool
+                json.tool.main()
+
+        os.unlink(infile + '.out')
+        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', '-i', infile]):
+                try:
+                    import json.tool
+                    json.tool.main()
+                except SystemExit:
+                    pass
+
+        self.assertEqual(opened, closed)
+        self.assertEqual(len(opened), 2)
+        self.assertEqual(len(opened), 2)
+
     def test_indent(self):
         input_ = '[1, 2]'
         expect = textwrap.dedent('''\
@@ -219,3 +279,20 @@ def test_broken_pipe_error(self):
         proc.stdout.close()
         proc.communicate(b'"{}"')
         self.assertEqual(proc.returncode, errno.EPIPE)
+
+
+def mock_open():
+    closed = []
+    opened = []
+    io_open = io.open
+
+    def _open(*args, **kwargs):
+        fd = io_open(*args, **kwargs)
+        opened.append(fd)
+        fd_close = fd.close
+        def close(self):
+            closed.append(self)
+            fd_close()
+        fd.close = types.MethodType(close, fd)
+        return fd
+    return closed, opened, _open

Misc/NEWS.d/next/Library/2018-06-27-14-46-03.bpo-33927.felCrI.rst

@@ -0,0 +1,2 @@
+``json.tool`` can now take the same file as input and ouput with the ``--in-place``
+flag. Patch by Rémi Lapeyre.