Patch #1630118: add a SpooledTemporaryFile class to tempfile.

1 files changed, 114 insertions, 0 deletions

diff --git a/Lib/tempfile.py b/Lib/tempfile.py
index ce03bb7..bb23785 100644
--- a/Lib/tempfile.py
+++ b/Lib/tempfile.py
@@ -19,6 +19,7 @@ This module also provides some data items to the user:
 
 __all__ = [
     "NamedTemporaryFile", "TemporaryFile", # high level safe interfaces
+    "SpooledTemporaryFile",
     "mkstemp", "mkdtemp",                  # low level safe interfaces
     "mktemp",                              # deprecated unsafe interface
     "TMP_MAX", "gettempprefix",            # constants
@@ -37,6 +38,11 @@ if _os.name == 'mac':
     import Carbon.Folders as _Folders
 
 try:
+  from cStringIO import StringIO as _StringIO
+except:
+  from StringIO import StringIO as _StringIO
+
+try:
     import fcntl as _fcntl
 except ImportError:
     def _set_cloexec(fd):
@@ -473,3 +479,111 @@ else:
         except:
             _os.close(fd)
             raise
+
+class SpooledTemporaryFile:
+    """Temporary file wrapper, specialized to switch from
+    StringIO to a real file when it exceeds a certain size or
+    when a fileno is needed.
+    """
+    _rolled = False
+
+    def __init__(self, max_size=0, mode='w+b', bufsize=-1,
+                 suffix="", prefix=template, dir=None):
+        self._file = _StringIO()
+        self._max_size = max_size
+        self._rolled = False
+        self._TemporaryFileArgs = (mode, bufsize, suffix, prefix, dir)
+
+    def _check(self, file):
+        if self._rolled: return
+        max_size = self._max_size
+        if max_size and file.tell() > max_size:
+            self.rollover()
+
+    def rollover(self):
+        if self._rolled: return
+        file = self._file
+        newfile = self._file = TemporaryFile(*self._TemporaryFileArgs)
+        del self._TemporaryFileArgs
+
+        newfile.write(file.getvalue())
+        newfile.seek(file.tell(), 0)
+
+        self._rolled = True
+        
+    # file protocol
+    def __iter__(self):
+        return self._file.__iter__()
+
+    def close(self):
+        self._file.close()
+
+    @property
+    def closed(self):
+        return self._file.closed
+
+    @property
+    def encoding(self):
+        return self._file.encoding
+
+    def fileno(self):
+        self.rollover()
+        return self._file.fileno()
+
+    def flush(self):
+        self._file.flush()
+
+    def isatty(self):
+        return self._file.isatty()
+
+    @property
+    def mode(self):
+        return self._file.mode
+
+    @property
+    def name(self):
+        return self._file.name
+
+    @property
+    def newlines(self):
+        return self._file.newlines
+
+    def next(self):
+        return self._file.next
+
+    def read(self, *args):
+        return self._file.read(*args)
+
+    def readline(self, *args):
+        return self._file.readline(*args)
+
+    def readlines(self, *args):
+        return self._file.readlines(*args)
+
+    def seek(self, *args):
+        self._file.seek(*args)
+
+    @property
+    def softspace(self):
+        return self._file.softspace
+
+    def tell(self):
+        return self._file.tell()
+
+    def truncate(self):
+        self._file.truncate()
+
+    def write(self, s):
+        file = self._file
+        rv = file.write(s)
+        self._check(file)
+        return rv
+
+    def writelines(self, iterable):
+        file = self._file
+        rv = file.writelines(iterable)
+        self._check(file)
+        return rv
+
+    def xreadlines(self, *args):
+        return self._file.xreadlines(*args)