diff --git a/Lib/tempfile.py b/Lib/tempfile.py index 0eb9ddeb6ac377..cadb0bed3cce3b 100644 --- a/Lib/tempfile.py +++ b/Lib/tempfile.py @@ -868,10 +868,14 @@ def write(self, s): return rv def writelines(self, iterable): - file = self._file - rv = file.writelines(iterable) - self._check(file) - return rv + if self._max_size == 0 or self._rolled: + return self._file.writelines(iterable) + + it = iter(iterable) + for line in it: + self.write(line) + if self._rolled: + return self._file.writelines(it) def detach(self): return self._file.detach() diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py index 7adc021d298254..d46d3c0f040601 100644 --- a/Lib/test/test_tempfile.py +++ b/Lib/test/test_tempfile.py @@ -1284,6 +1284,34 @@ def test_writelines(self): buf = f.read() self.assertEqual(buf, b'xyz') + def test_writelines_rollover(self): + # Verify writelines rolls over before exhausting the iterator + f = self.do_create(max_size=2) + + def it(): + yield b'xy' + self.assertFalse(f._rolled) + yield b'z' + self.assertTrue(f._rolled) + + f.writelines(it()) + pos = f.seek(0) + self.assertEqual(pos, 0) + buf = f.read() + self.assertEqual(buf, b'xyz') + + def test_writelines_fast_path(self): + f = self.do_create(max_size=2) + f.write(b'abc') + self.assertTrue(f._rolled) + + f.writelines([b'd', b'e', b'f']) + pos = f.seek(0) + self.assertEqual(pos, 0) + buf = f.read() + self.assertEqual(buf, b'abcdef') + + def test_writelines_sequential(self): # A SpooledTemporaryFile should hold exactly max_size bytes, and roll # over afterward diff --git a/Misc/NEWS.d/next/Security/2024-11-28-20-29-21.gh-issue-127371.PeEhUd.rst b/Misc/NEWS.d/next/Security/2024-11-28-20-29-21.gh-issue-127371.PeEhUd.rst new file mode 100644 index 00000000000000..029c348918e0d1 --- /dev/null +++ b/Misc/NEWS.d/next/Security/2024-11-28-20-29-21.gh-issue-127371.PeEhUd.rst @@ -0,0 +1,3 @@ +Avoid unbounded buffering for :meth:`!tempfile.SpooledTemporaryFile.writelines`. +Previously, disk spillover was only checked after the lines iterator had been +exhausted. This is now done after each line is written.