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