diff --git a/Lib/json/tool.py b/Lib/json/tool.py index 0490b8c0be11dfe..a0a8233ec61b501 100644 --- a/Lib/json/tool.py +++ b/Lib/json/tool.py @@ -69,6 +69,10 @@ def main(): if options.outfile is None: out = sys.stdout else: + if (options.outfile.exists() and + options.outfile.expanduser().samefile(infile.name)): + # exhaust the generator before opening for writing + objs = list(objs) out = options.outfile.open('w', encoding='utf-8') with out as outfile: for obj in objs: diff --git a/Lib/test/support/script_helper.py b/Lib/test/support/script_helper.py index 6d699c8486cd20d..10b04950a82ec48 100644 --- a/Lib/test/support/script_helper.py +++ b/Lib/test/support/script_helper.py @@ -2,6 +2,7 @@ # e.g. test_cmd_line, test_cmd_line_script and test_runpy import collections +import contextlib import importlib import sys import os @@ -196,9 +197,12 @@ def spawn_python(*args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, **kw): def kill_python(p): """Run the given Popen process until completion and return stdout.""" - p.stdin.close() - data = p.stdout.read() - p.stdout.close() + data = None + if p.stdin: + p.stdin.close() + if p.stdout: + data = p.stdout.read() + p.stdout.close() # try to cleanup the child so we don't appear to leak when running # with regrtest -R. p.wait() @@ -206,6 +210,16 @@ def kill_python(p): return data +@contextlib.contextmanager +def killing(process): + """Context manager that kills and cleans up the given Popen on exit.""" + try: + yield process + finally: + process.kill() + kill_python(process) + + def make_script(script_dir, script_basename, source, omit_suffix=False): script_filename = script_basename if not omit_suffix: diff --git a/Lib/test/test_json/test_tool.py b/Lib/test/test_json/test_tool.py index 1d7fca6efb1cc72..c67f9115a93479a 100644 --- a/Lib/test/test_json/test_tool.py +++ b/Lib/test/test_json/test_tool.py @@ -7,7 +7,7 @@ from test import support from test.support import os_helper -from test.support.script_helper import assert_python_ok +from test.support.script_helper import assert_python_ok, killing class TestTool(unittest.TestCase): @@ -146,6 +146,40 @@ def test_jsonlines(self): self.assertEqual(process.stdout, self.jsonlines_expect) self.assertEqual(process.stderr, '') + def test_writing_jsonlines_in_place(self): + infile = self._create_infile(data=self.jsonlines_raw) + rc, out, err = assert_python_ok('-m', 'json.tool', '--json-lines', infile, infile) + with open(infile, "r", encoding="utf-8") as fp: + self.assertEqual(fp.read(), self.jsonlines_expect) + self.assertEqual(rc, 0) + self.assertEqual(out, b'') + self.assertEqual(err, b'') + + def test_jsonlines_from_a_stream(self): + stream_code = textwrap.dedent(""" + # print a stream of JSON integers + i = 0 + while True: + print(i) + i += 1 + """) + stream = subprocess.Popen((sys.executable, '-u', '-c', stream_code), + stdout=subprocess.PIPE, + bufsize=0, + ) + args = sys.executable, '-u', '-m', 'json.tool', '--json-lines' + with killing(stream), \ + self.assertRaises(subprocess.TimeoutExpired) as timeout: + subprocess.run(args, stdin=stream.stdout, + capture_output=True, timeout=1, + bufsize=0) + + self.assertIsNotNone(timeout.exception.stdout) + self.assertEqual( + timeout.exception.stdout.splitlines()[:3], + [b'0', b'1', b'2'] + ) + def test_help_flag(self): rc, out, err = assert_python_ok('-m', 'json.tool', '-h') self.assertEqual(rc, 0)