From f17fc9acfe2aee7c8daa38eefeea64685ce4b6df Mon Sep 17 00:00:00 2001
From: Oliver Beckstein <orbeckst@gmail.com>
Date: Mon, 9 Jun 2014 11:58:13 -0700
Subject: [PATCH] file open/close functions can process any stream *filename*

---
 propka/lib.py | 43 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 37 insertions(+), 6 deletions(-)

diff --git a/propka/lib.py b/propka/lib.py
index 0bbb0a6..caa62ff 100644
--- a/propka/lib.py
+++ b/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:
+            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:
-        res = open(filename,'w')
+        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()