diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/tempfile.py | 9 | ||||
| -rw-r--r-- | Lib/test/test_tempfile.py | 21 |
2 files changed, 28 insertions, 2 deletions
diff --git a/Lib/tempfile.py b/Lib/tempfile.py index 3ec6b4a..e3afa3b 100644 --- a/Lib/tempfile.py +++ b/Lib/tempfile.py @@ -578,8 +578,13 @@ class SpooledTemporaryFile: def tell(self): return self._file.tell() - def truncate(self): - self._file.truncate() + def truncate(self, size=None): + if size is None: + self._file.truncate() + else: + if size > self._max_size: + self.rollover() + self._file.truncate(size) def write(self, s): file = self._file diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py index 014fca0..5b0bca7 100644 --- a/Lib/test/test_tempfile.py +++ b/Lib/test/test_tempfile.py @@ -846,6 +846,27 @@ class test_SpooledTemporaryFile(TC): pass self.assertRaises(ValueError, use_closed) + def test_truncate_with_size_parameter(self): + # A SpooledTemporaryFile can be truncated to zero size + f = tempfile.SpooledTemporaryFile(max_size=10) + f.write(b'abcdefg\n') + f.seek(0) + f.truncate() + self.assertFalse(f._rolled) + self.assertEqual(f._file.getvalue(), b'') + # A SpooledTemporaryFile can be truncated to a specific size + f = tempfile.SpooledTemporaryFile(max_size=10) + f.write(b'abcdefg\n') + f.truncate(4) + self.assertFalse(f._rolled) + self.assertEqual(f._file.getvalue(), b'abcd') + # A SpooledTemporaryFile rolls over if truncated to large size + f = tempfile.SpooledTemporaryFile(max_size=10) + f.write(b'abcdefg\n') + f.truncate(20) + self.assertTrue(f._rolled) + if has_stat: + self.assertEqual(os.fstat(f.fileno()).st_size, 20) test_classes.append(test_SpooledTemporaryFile) |