/usr/share/pyshared/restkit/tee.py is in python-restkit 4.2.2-1.
This file is owned by root:root, with mode 0o644.
The actual contents of the file can be viewed below.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | # -*- coding: utf-8 -
#
# This file is part of restkit released under the MIT license.
# See the NOTICE for more information.
"""
TeeInput replace old FileInput. It use a file
if size > MAX_BODY or memory. It's now possible to rewind
read or restart etc ... It's based on TeeInput from Gunicorn.
"""
import copy
import os
try:
from cStringIO import StringIO
except ImportError:
from StringIO import StringIO
import tempfile
from restkit import conn
class TeeInput(object):
CHUNK_SIZE = conn.CHUNK_SIZE
def __init__(self, stream):
self.buf = StringIO()
self.eof = False
if isinstance(stream, basestring):
stream = StringIO(stream)
self.tmp = StringIO()
else:
self.tmp = tempfile.TemporaryFile()
self.stream = stream
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, traceback):
return
def seek(self, offset, whence=0):
""" naive implementation of seek """
current_size = self._tmp_size()
diff = 0
if whence == 0:
diff = offset - current_size
elif whence == 2:
diff = (self.tmp.tell() + offset) - current_size
elif whence == 3 and not self.eof:
# we read until the end
while True:
self.tmp.seek(0, 2)
if not self._tee(self.CHUNK_SIZE):
break
if not self.eof and diff > 0:
self._ensure_length(StringIO(), diff)
self.tmp.seek(offset, whence)
def flush(self):
self.tmp.flush()
def read(self, length=-1):
""" read """
if self.eof:
return self.tmp.read(length)
if length < 0:
buf = StringIO()
buf.write(self.tmp.read())
while True:
chunk = self._tee(self.CHUNK_SIZE)
if not chunk:
break
buf.write(chunk)
return buf.getvalue()
else:
dest = StringIO()
diff = self._tmp_size() - self.tmp.tell()
if not diff:
dest.write(self._tee(length))
return self._ensure_length(dest, length)
else:
l = min(diff, length)
dest.write(self.tmp.read(l))
return self._ensure_length(dest, length)
def readline(self, size=-1):
if self.eof:
return self.tmp.readline()
orig_size = self._tmp_size()
if self.tmp.tell() == orig_size:
if not self._tee(self.CHUNK_SIZE):
return ''
self.tmp.seek(orig_size)
# now we can get line
line = self.tmp.readline()
if line.find("\n") >=0:
return line
buf = StringIO()
buf.write(line)
while True:
orig_size = self.tmp.tell()
data = self._tee(self.CHUNK_SIZE)
if not data:
break
self.tmp.seek(orig_size)
buf.write(self.tmp.readline())
if data.find("\n") >= 0:
break
return buf.getvalue()
def readlines(self, sizehint=0):
total = 0
lines = []
line = self.readline()
while line:
lines.append(line)
total += len(line)
if 0 < sizehint <= total:
break
line = self.readline()
return lines
def close(self):
if not self.eof:
# we didn't read until the end
self._close_unreader()
return self.tmp.close()
def next(self):
r = self.readline()
if not r:
raise StopIteration
return r
__next__ = next
def __iter__(self):
return self
def _tee(self, length):
""" fetch partial body"""
buf2 = self.buf
buf2.seek(0, 2)
chunk = self.stream.read(length)
if chunk:
self.tmp.write(chunk)
self.tmp.flush()
self.tmp.seek(0, 2)
return chunk
self._finalize()
return ""
def _finalize(self):
""" here we wil fetch final trailers
if any."""
self.eof = True
def _tmp_size(self):
if hasattr(self.tmp, 'fileno'):
return int(os.fstat(self.tmp.fileno())[6])
else:
return len(self.tmp.getvalue())
def _ensure_length(self, dest, length):
if len(dest.getvalue()) < length:
data = self._tee(length - len(dest.getvalue()))
dest.write(data)
return dest.getvalue()
class ResponseTeeInput(TeeInput):
CHUNK_SIZE = conn.CHUNK_SIZE
def __init__(self, resp, connection, should_close=False):
self.buf = StringIO()
self.resp = resp
self.stream =resp.body_stream()
self.connection = connection
self.should_close = should_close
self.eof = False
# set temporary body
clen = int(resp.headers.get('content-length') or -1)
if clen >= 0:
if (clen <= conn.MAX_BODY):
self.tmp = StringIO()
else:
self.tmp = tempfile.TemporaryFile()
else:
self.tmp = tempfile.TemporaryFile()
def close(self):
if not self.eof:
# we didn't read until the end
self._close_unreader()
return self.tmp.close()
def _close_unreader(self):
if not self.eof:
self.stream.close()
self.connection.release(self.should_close)
def _finalize(self):
""" here we wil fetch final trailers
if any."""
self.eof = True
self._close_unreader()
|