Package restkit :: Module tee
[hide private]
[frames] | no frames]

Source Code for Module restkit.tee

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