file open/close functions can process any stream *filename*

propka/lib.py

@@ -9,18 +9,49 @@ import pkg_resources
 # file I/O
 #
 def open_file_for_reading(filename):
+    """Open file or file-like stream  *filename* for reading.
+
+    *filename* may be a string and then it is opened but if it is a
+    file-like object (such as an open :class:`file` or
+    :class:`StringIO.StringIO` --- really anything with ``next()``,
+    ``read()``, ``readlines()``, ``readline``, ``close`` methods) then
+    the object is just passed through (the stream is attempted to be
+    reset to the beginning with ``fseek(0)``).
+    """
+    if hasattr(filename, 'next') and hasattr(filename, 'read') \
+            and hasattr(filename, 'readline') and hasattr(filename, 'readlines') \
+            and hasattr(filename, 'close'):
+        # already a stream
+        try:
+            filename.fseek(0)
+        except AttributeError:
+            pass
+        return filename
+
     try:
         f = open(filename,'r')
     except:
-        raise Exception('Cannot find file %s' %filename)
+        raise IOError('Cannot find file %s' %filename)
     return f
 
 def open_file_for_writing(filename):
+    """Open file or file-like stream for writing"""
+    if hasattr(filename, 'write') and hasattr(filename, 'writeline') and hasattr(filename, 'writelines') \
+            and hasattr(filename, 'close'):
+        # already a stream
         try:
-        res = open(filename,'w')
+            mode = filename.mode
+        except AttributeError:
+            mode = "w"
+        if not ("w" in mode or "a" in mode or "+" in mode):
+            raise IOError("File/stream not open for writing")
+        return filename
+
+    try:
+        f = open(filename,'w')
     except:
         raise Exception('Could not open %s'%filename)
-    return res
+    return f
 
 #
 # bookkeeping etc.
@@ -193,10 +224,10 @@ def writeFile(filename, lines):
     """
     Writes a new file
     """
-    file = open(filename, 'w')
+    f = open_file_for_writing(filename)
 
     for line in lines:
-        file.write( "%s\n" % (line) )
-    file.close()
+        f.write( "%s\n" % (line) )
+    f.close()