diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/fileinput.py | 6 | ||||
| -rw-r--r-- | Lib/test/test_fileinput.py | 24 |
2 files changed, 30 insertions, 0 deletions
diff --git a/Lib/fileinput.py b/Lib/fileinput.py index 90a600b..a25a021 100644 --- a/Lib/fileinput.py +++ b/Lib/fileinput.py @@ -238,6 +238,12 @@ class FileInput: self.nextfile() self._files = () + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + self.close() + def __iter__(self): return self diff --git a/Lib/test/test_fileinput.py b/Lib/test/test_fileinput.py index aae16ea..f312882 100644 --- a/Lib/test/test_fileinput.py +++ b/Lib/test/test_fileinput.py @@ -231,6 +231,30 @@ class FileInputTests(unittest.TestCase): ## finally: ## remove_tempfiles(t1) + def test_context_manager(self): + try: + t1 = writeTmp(1, ["A\nB\nC"]) + t2 = writeTmp(2, ["D\nE\nF"]) + with FileInput(files=(t1, t2)) as fi: + lines = list(fi) + self.assertEqual(lines, ["A\n", "B\n", "C", "D\n", "E\n", "F"]) + self.assertEqual(fi.filelineno(), 3) + self.assertEqual(fi.lineno(), 6) + self.assertEqual(fi._files, ()) + finally: + remove_tempfiles(t1, t2) + + def test_close_on_exception(self): + try: + t1 = writeTmp(1, [""]) + with FileInput(files=t1) as fi: + raise IOError + except IOError: + self.assertEqual(fi._files, ()) + finally: + remove_tempfiles(t1) + + def test_main(): run_unittest(BufferSizesTests, FileInputTests) |