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