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  from StringIO import StringIO 
 15  import tempfile 
 16   
 17  from restkit import sock 
18 19 -class TeeInput(object):
20 21 CHUNK_SIZE = sock.CHUNK_SIZE 22
23 - def __init__(self, socket, parser, buf, maybe_close=None):
24 self.buf = StringIO() 25 self.parser = parser 26 self._sock = socket 27 self.maybe_close = maybe_close 28 self._is_socket = True 29 self._len = parser.content_len 30 31 if self._len and self._len < sock.MAX_BODY: 32 self.tmp = StringIO() 33 else: 34 self.tmp = tempfile.TemporaryFile() 35 36 if buf.len > 0: 37 chunk, self.buf = parser.filter_body(buf) 38 if chunk: 39 self.tmp.write(chunk) 40 self.tmp.flush() 41 self._finalize() 42 self.tmp.seek(0) 43 del buf
44 45 @property
46 - def len(self):
47 if self._len: return self._len 48 49 if self._is_socket: 50 pos = self.tmp.tell() 51 self.tmp.seek(0, 2) 52 while True: 53 if not self._tee(self.CHUNK_SIZE): 54 break 55 self.tmp.seek(pos) 56 self._len = self._tmp_size() 57 return self._len
58
59 - def seek(self, offset, whence=0):
60 """ naive implementation of seek """ 61 if self._is_socket: 62 self.tmp.seek(0, 2) 63 while True: 64 if not self._tee(self.CHUNK_SIZE): 65 break 66 self.tmp.seek(offset, whence)
67
68 - def flush(self):
69 self.tmp.flush()
70
71 - def read(self, length=-1):
72 """ read """ 73 if not self._is_socket: 74 return self.tmp.read(length) 75 76 if length < 0: 77 buf = StringIO() 78 buf.write(self.tmp.read()) 79 while True: 80 chunk = self._tee(self.CHUNK_SIZE) 81 if not chunk: 82 break 83 buf.write(chunk) 84 return buf.getvalue() 85 else: 86 dest = StringIO() 87 diff = self._tmp_size() - self.tmp.tell() 88 if not diff: 89 dest.write(self._tee(length)) 90 return self._ensure_length(dest, length) 91 else: 92 l = min(diff, length) 93 dest.write(self.tmp.read(l)) 94 return self._ensure_length(dest, length)
95
96 - def readline(self, size=-1):
97 if not self._is_socket: 98 return self.tmp.readline() 99 100 orig_size = self._tmp_size() 101 if self.tmp.tell() == orig_size: 102 if not self._tee(self.CHUNK_SIZE): 103 return '' 104 self.tmp.seek(orig_size) 105 106 # now we can get line 107 line = self.tmp.readline() 108 if line.find("\n") >=0: 109 return line 110 111 buf = StringIO() 112 buf.write(line) 113 while True: 114 orig_size = self.tmp.tell() 115 data = self._tee(self.CHUNK_SIZE) 116 if not data: 117 break 118 self.tmp.seek(orig_size) 119 buf.write(self.tmp.readline()) 120 if data.find("\n") >= 0: 121 break 122 return buf.getvalue()
123
124 - def readlines(self, sizehint=0):
125 total = 0 126 lines = [] 127 line = self.readline() 128 while line: 129 lines.append(line) 130 total += len(line) 131 if 0 < sizehint <= total: 132 break 133 line = self.readline() 134 return lines
135
136 - def next(self):
137 r = self.readline() 138 if not r: 139 raise StopIteration 140 return r
141 __next__ = next 142
143 - def __iter__(self):
144 return self
145
146 - def _tee(self, length):
147 """ fetch partial body""" 148 buf2 = self.buf 149 buf2.seek(0, 2) 150 while True: 151 chunk, buf2 = self.parser.filter_body(buf2) 152 if chunk: 153 self.tmp.write(chunk) 154 self.tmp.flush() 155 self.tmp.seek(0, 2) 156 self.buf = StringIO() 157 self.buf.write(buf2.getvalue()) 158 return chunk 159 160 if self.parser.body_eof(): 161 break 162 163 data = self._sock.recv(length) 164 buf2.write(data) 165 166 self._finalize() 167 return ""
168
169 - def _finalize(self):
170 """ here we wil fetch final trailers 171 if any.""" 172 173 if self.parser.body_eof(): 174 if callable(self.maybe_close): 175 self.maybe_close() 176 177 self.buf = StringIO() 178 self._is_socket = False
179
180 - def _tmp_size(self):
181 if isinstance(self.tmp, StringIO): 182 return self.tmp.len 183 else: 184 return int(os.fstat(self.tmp.fileno())[6])
185
186 - def _ensure_length(self, dest, length):
187 if not dest.len or not self._len: 188 return dest.getvalue() 189 while True: 190 if dest.len >= length: 191 break 192 data = self._tee(length - dest.len) 193 if not data: 194 break 195 dest.write(data) 196 return dest.getvalue()
197