Patch #1215184: FileInput now can be given an opening hook which can

birkenfeld · birkenfeld · commit c98eeede175a · 2006-02-19T14:57:47.000Z
be used to control how files are opened.
diff --git a/Doc/lib/libfileinput.tex b/Doc/lib/libfileinput.tex
@@ -43,17 +43,23 @@ \section{\module{fileinput} ---
 character; lines are returned including the trailing newline when it
 is present.
 
+You can control how files are opened by providing an opening hook via the
+\var{openhook} parameter to \function{input()} or \class{FileInput()}.
+The hook must be a function that takes two arguments, \var{filename}
+and \var{mode}, and returns an accordingly opened file-like object.
+Two useful hooks are already provided by this module.
+
 The following function is the primary interface of this module:
 
-\begin{funcdesc}{input}{\optional{files\optional{,
-                       inplace\optional{, backup\optional{, mode}}}}}
+\begin{funcdesc}{input}{\optional{files\optional{, inplace\optional{,
+                        backup\optional{, mode\optional{, openhook}}}}}}
   Create an instance of the \class{FileInput} class.  The instance
   will be used as global state for the functions of this module, and
   is also returned to use during iteration.  The parameters to this
   function will be passed along to the constructor of the
   \class{FileInput} class.
 
-  \versionchanged[Added the \var{mode} parameter]{2.5}
+  \versionchanged[Added the \var{mode} and \var{openhook} parameters]{2.5}
 \end{funcdesc}
 
 
@@ -115,7 +121,8 @@ \section{\module{fileinput} ---
 module is available for subclassing as well:
 
 \begin{classdesc}{FileInput}{\optional{files\optional{,
-                             inplace\optional{, backup\optional{, mode}}}}}
+                             inplace\optional{, backup\optional{,
+                             mode\optional{, openhook}}}}}}
   Class \class{FileInput} is the implementation; its methods
   \method{filename()}, \method{fileno()}, \method{lineno()},
   \method{fileline()}, \method{isfirstline()}, \method{isstdin()},
@@ -131,7 +138,12 @@ \section{\module{fileinput} ---
   \function{open()}. It must be one of \code{'r'}, \code{'rU'},
   \code{'U'} and \code{'rb'}.
 
-  \versionchanged[Added the \var{mode} parameter]{2.5}
+  The \var{openhook}, when given, must be a function that takes two arguments,
+  \var{filename} and \var{mode}, and returns an accordingly opened
+  file-like object.
+  You cannot use \var{inplace} and \var{openhook} together.
+
+  \versionchanged[Added the \var{mode} and \var{openhook} parameters]{2.5}
 \end{classdesc}
 
 \strong{Optional in-place filtering:} if the keyword argument
@@ -148,3 +160,29 @@ \section{\module{fileinput} ---
 
 \strong{Caveat:} The current implementation does not work for MS-DOS
 8+3 filesystems.
+
+
+The two following opening hooks are provided by this module:
+
+\begin{funcdesc}{hook_compressed}{filename, mode}
+  Transparently opens files compressed with gzip and bzip2 using
+  the \module{gzip} and \module{bz2} modules.
+
+  Usage example: 
+  \samp{fi = fileinput.FileInput(openhook=fileinput.hook_compressed)}
+
+  \versionadded{2.5}
+\end{funcdesc}
+
+\begin{funcdesc}{hook_encoded}{encoding}
+  Returns a hook which opens each file with \function{codecs.open()},
+  using the given \var{encoding} to read the file.
+
+  Usage example:
+  \samp{fi = fileinput.FileInput(openhook=fileinput.hook_encoded("iso-8859-1"))}
+
+  \note{With this hook, \class{FileInput} might return Unicode strings
+        depending on the specified \var{encoding}.}
+  \versionadded{2.5}
+\end{funcdesc}
+
diff --git a/Lib/fileinput.py b/Lib/fileinput.py
@@ -88,8 +88,9 @@
 
 DEFAULT_BUFSIZE = 8*1024
 
-def input(files=None, inplace=0, backup="", bufsize=0, mode="r"):
-    """input([files[, inplace[, backup[, mode]]]])
+def input(files=None, inplace=0, backup="", bufsize=0, 
+          mode="r", openhook=None):
+    """input([files[, inplace[, backup[, mode[, openhook]]]]])
 
     Create an instance of the FileInput class. The instance will be used
     as global state for the functions of this module, and is also returned
@@ -99,7 +100,7 @@ def input(files=None, inplace=0, backup="", bufsize=0, mode="r"):
     global _state
     if _state and _state._file:
         raise RuntimeError, "input() already active"
-    _state = FileInput(files, inplace, backup, bufsize, mode)
+    _state = FileInput(files, inplace, backup, bufsize, mode, openhook)
     return _state
 
 def close():
@@ -181,7 +182,7 @@ def isstdin():
     return _state.isstdin()
 
 class FileInput:
-    """class FileInput([files[, inplace[, backup[, mode]]]])
+    """class FileInput([files[, inplace[, backup[, mode[, openhook]]]]])
 
     Class FileInput is the implementation of the module; its methods
     filename(), lineno(), fileline(), isfirstline(), isstdin(), fileno(),
@@ -193,7 +194,8 @@ class FileInput:
     sequential order; random access and readline() cannot be mixed.
     """
 
-    def __init__(self, files=None, inplace=0, backup="", bufsize=0, mode="r"):
+    def __init__(self, files=None, inplace=0, backup="", bufsize=0, 
+                 mode="r", openhook=None):
         if isinstance(files, basestring):
             files = (files,)
         else:
@@ -222,6 +224,11 @@ def __init__(self, files=None, inplace=0, backup="", bufsize=0, mode="r"):
             raise ValueError("FileInput opening mode must be one of "
                              "'r', 'rU', 'U' and 'rb'")
         self._mode = mode
+        if inplace and openhook:
+            raise ValueError("FileInput cannot use an opening hook in inplace mode")
+        elif openhook and not callable(openhook):
+            raise ValueError("FileInput openhook must be callable")
+        self._openhook = openhook
 
     def __del__(self):
         self.close()
@@ -332,7 +339,10 @@ def readline(self):
                     sys.stdout = self._output
                 else:
                     # This may raise IOError
-                    self._file = open(self._filename, self._mode)
+                    if self._openhook:
+                        self._file = self._openhook(self._filename, self._mode)
+                    else:
+                        self._file = open(self._filename, self._mode)
         self._buffer = self._file.readlines(self._bufsize)
         self._bufindex = 0
         if not self._buffer:
@@ -364,6 +374,26 @@ def isfirstline(self):
     def isstdin(self):
         return self._isstdin
 
+
+def hook_compressed(filename, mode):
+    ext = os.path.splitext(filename)[1]
+    if ext == '.gz':
+        import gzip
+        return gzip.open(filename, mode)
+    elif ext == '.bz2':
+        import bz2
+        return bz2.BZ2File(filename, mode)
+    else:
+        return open(filename, mode)
+
+
+def hook_encoded(encoding):
+    import codecs
+    def openhook(filename, mode):
+        return codecs.open(filename, mode, encoding)
+    return openhook
+
+
 def _test():
     import getopt
     inplace = 0
diff --git a/Lib/test/test_fileinput.py b/Lib/test/test_fileinput.py
@@ -6,7 +6,7 @@
 from test.test_support import verify, verbose, TESTFN, TestFailed
 import sys, os, re
 from StringIO import StringIO
-from fileinput import FileInput
+from fileinput import FileInput, hook_encoded
 
 # The fileinput module has 2 interfaces: the FileInput class which does
 # all the work, and a few functions (input, etc.) that use a global _state
@@ -200,3 +200,25 @@ def writeFiles():
     verify(lines == ["A\n", "B\n", "C\n", "D"])
 finally:
     remove_tempfiles(t1)
+
+if verbose:
+    print "18. Test file opening hook"
+try:
+    # cannot use openhook and inplace mode
+    fi = FileInput(inplace=1, openhook=lambda f,m: None)
+    raise TestFailed("FileInput should raise if both inplace "
+                     "and openhook arguments are given")
+except ValueError:
+    pass
+try:
+    fi = FileInput(openhook=1)
+    raise TestFailed("FileInput should check openhook for being callable")
+except ValueError:
+    pass
+try:
+    t1 = writeTmp(1, ["A\nB"])
+    fi = FileInput(files=t1, openhook=hook_encoded("rot13"))
+    lines = list(fi)
+    verify(lines == ["N\n", "O"])
+finally:
+    remove_tempfiles(t1)
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -366,6 +366,9 @@ Extension Modules
 Library
 -------
 
+- Patch #1215184: FileInput now can be given an opening hook which can
+  be used to control how files are opened.
+
 - Patch #1212287: fileinput.input() now has a mode parameter for
   specifying the file mode input files should be opened with.
 
