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