#1286: allow using fileinput.FileInput as context manager.

birkenfeld · birkenfeld · commit 6cb7b6593ee0 · 2010-07-31T20:08:15.000Z
diff --git a/Doc/library/fileinput.rst b/Doc/library/fileinput.rst
@@ -24,7 +24,7 @@ as the first argument to :func:`.input`.  A single file name is also allowed.
 
 All files are opened in text mode by default, but you can override this by
 specifying the *mode* parameter in the call to :func:`.input` or
-:class:`FileInput()`.  If an I/O error occurs during opening or reading a file,
+:class:`FileInput`.  If an I/O error occurs during opening or reading a file,
 :exc:`IOError` is raised.
 
 If ``sys.stdin`` is used more than once, the second and further use will return
@@ -54,6 +54,16 @@ The following function is the primary interface of this module:
    during iteration.  The parameters to this function will be passed along to the
    constructor of the :class:`FileInput` class.
 
+   The :class:`FileInput` instance can be used as a context manager in the
+   :keyword:`with` statement.  In this example, *input* is closed after the
+   :keyword:`with` statement is exited, even if an exception occurs::
+
+      with fileinput.input(files=('spam.txt', 'eggs.txt')) as input:
+          process(input)
+
+   .. versionchanged:: 3.2
+      Can be used as a context manager.
+
 
 The following functions use the global state created by :func:`fileinput.input`;
 if there is no active state, :exc:`RuntimeError` is raised.
@@ -132,13 +142,23 @@ available for subclassing as well:
    *filename* and *mode*, and returns an accordingly opened file-like object. You
    cannot use *inplace* and *openhook* together.
 
+   A :class:`FileInput` instance can be used as a context manager in the
+   :keyword:`with` statement.  In this example, *input* is closed after the
+   :keyword:`with` statement is exited, even if an exception occurs::
+
+      with FileInput(files=('spam.txt', 'eggs.txt')) as input:
+          process(input)
+
+   .. versionchanged:: 3.2
+      Can be used as a context manager.
+
 
-**Optional in-place filtering:** if the keyword argument ``inplace=1`` is passed
-to :func:`fileinput.input` or to the :class:`FileInput` constructor, the file is
-moved to a backup file and standard output is directed to the input file (if a
-file of the same name as the backup file already exists, it will be replaced
-silently).  This makes it possible to write a filter that rewrites its input
-file in place.  If the *backup* parameter is given (typically as
+**Optional in-place filtering:** if the keyword argument ``inplace=True`` is
+passed to :func:`fileinput.input` or to the :class:`FileInput` constructor, the
+file is moved to a backup file and standard output is directed to the input file
+(if a file of the same name as the backup file already exists, it will be
+replaced silently).  This makes it possible to write a filter that rewrites its
+input file in place.  If the *backup* parameter is given (typically as
 ``backup='.<some extension>'``), it specifies the extension for the backup file,
 and the backup file remains around; by default, the extension is ``'.bak'`` and
 it is deleted when the output file is closed.  In-place filtering is disabled
diff --git a/Lib/fileinput.py b/Lib/fileinput.py
@@ -238,6 +238,12 @@ def close(self):
         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
@@ -231,6 +231,30 @@ def test_file_opening_hook(self):
 ##         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)
 
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -15,6 +15,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #1286: Allow using fileinput.FileInput as a context manager.
+
 - Add lfu_cache() and lru_cache() decorators to the functools module.
 
 
