Package restkit :: Package http :: Module body
[hide private]
[frames] | no frames]

Source Code for Module restkit.http.body

  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  import sys 
  7  import zlib 
  8   
  9  try: 
 10      from cStringIO import StringIO 
 11  except ImportError: 
 12      from StringIO import StringIO 
 13   
 14  from restkit.errors import NoMoreData, ChunkMissingTerminator, \ 
 15  InvalidChunkSize 
 16   
17 -class ChunkedReader(object):
18 - def __init__(self, req, unreader):
19 self.unreader = unreader 20 self.req = req 21 self.parser = self.parse_chunked(unreader) 22 self.buf = StringIO()
23
24 - def read(self, size):
25 if not isinstance(size, (int, long)): 26 raise TypeError("size must be an integral type") 27 if size <= 0: 28 raise ValueError("Size must be positive.") 29 if size == 0: 30 return "" 31 32 if self.parser: 33 while self.buf.tell() < size: 34 try: 35 self.buf.write(self.parser.next()) 36 except StopIteration: 37 self.parser = None 38 break 39 40 data = self.buf.getvalue() 41 ret, rest = data[:size], data[size:] 42 self.buf.truncate(0) 43 self.buf.write(rest) 44 return ret
45
46 - def parse_trailers(self, unreader, data, eof=False):
47 buf = StringIO() 48 buf.write(data) 49 50 idx = buf.getvalue().find("\r\n\r\n") 51 done = buf.getvalue()[:2] == "\r\n" 52 53 while idx < 0 and not done: 54 self.get_data(unreader, buf) 55 idx = buf.getvalue().find("\r\n\r\n") 56 done = buf.getvalue()[:2] == "\r\n" 57 if done: 58 unreader.unread(buf.getvalue()[2:]) 59 return "" 60 self.req.trailers = self.req.parse_headers(buf.getvalue()[:idx]) 61 unreader.unread(buf.getvalue()[idx+4:])
62
63 - def parse_chunked(self, unreader):
64 (size, rest) = self.parse_chunk_size(unreader) 65 while size > 0: 66 while size > len(rest): 67 size -= len(rest) 68 yield rest 69 rest = unreader.read() 70 if not rest: 71 raise NoMoreData() 72 yield rest[:size] 73 # Remove \r\n after chunk 74 rest = rest[size:] 75 while len(rest) < 2: 76 rest += unreader.read() 77 if rest[:2] != '\r\n': 78 raise ChunkMissingTerminator(rest[:2]) 79 (size, rest) = self.parse_chunk_size(unreader, data=rest[2:])
80
81 - def parse_chunk_size(self, unreader, data=None):
82 buf = StringIO() 83 if data is not None: 84 buf.write(data) 85 86 idx = buf.getvalue().find("\r\n") 87 while idx < 0: 88 self.get_data(unreader, buf) 89 idx = buf.getvalue().find("\r\n") 90 91 data = buf.getvalue() 92 line, rest_chunk = data[:idx], data[idx+2:] 93 94 chunk_size = line.split(";", 1)[0].strip() 95 try: 96 chunk_size = int(chunk_size, 16) 97 except ValueError: 98 raise InvalidChunkSize(chunk_size) 99 100 if chunk_size == 0: 101 try: 102 self.parse_trailers(unreader, rest_chunk) 103 except NoMoreData: 104 pass 105 return (0, None) 106 return (chunk_size, rest_chunk)
107
108 - def get_data(self, unreader, buf):
109 data = unreader.read() 110 if not data: 111 raise NoMoreData() 112 buf.write(data)
113
114 -class LengthReader(object):
115 - def __init__(self, req, unreader, length):
116 self.req = req 117 self.unreader = unreader 118 self.length = length
119
120 - def read(self, size):
121 if not isinstance(size, (int, long)): 122 raise TypeError("size must be an integral type") 123 124 size = min(self.length, size) 125 if size < 0: 126 raise ValueError("Size must be positive.") 127 if size == 0: 128 return "" 129 130 buf = StringIO() 131 data = self.unreader.read() 132 while data: 133 buf.write(data) 134 if buf.tell() >= size: 135 break 136 data = self.unreader.read() 137 138 139 buf = buf.getvalue() 140 ret, rest = buf[:size], buf[size:] 141 self.unreader.unread(rest) 142 self.length -= size 143 return ret
144
145 -class EOFReader(object):
146 - def __init__(self, req, unreader):
147 self.req = req 148 self.unreader = unreader 149 self.buf = StringIO() 150 self.finished = False
151
152 - def read(self, size):
153 if not isinstance(size, (int, long)): 154 raise TypeError("size must be an integral type") 155 if size < 0: 156 raise ValueError("Size must be positive.") 157 if size == 0: 158 return "" 159 160 if self.finished: 161 data = self.buf.getvalue() 162 ret, rest = data[:size], data[size:] 163 self.buf.truncate(0) 164 self.buf.write(rest) 165 return ret 166 167 data = self.unreader.read() 168 while data: 169 self.buf.write(data) 170 if self.buf.tell() > size: 171 break 172 data = self.unreader.read() 173 174 if not data: 175 self.finished = True 176 177 data = self.buf.getvalue() 178 ret, rest = data[:size], data[size:] 179 self.buf.truncate(0) 180 self.buf.write(rest) 181 return ret
182
183 -class Body(object):
184 - def __init__(self, reader):
185 self.reader = reader 186 self.req = reader.req 187 self.buf = StringIO() 188 self.closed = False
189
190 - def close(self):
191 """ Close the socket if needed """ 192 if self.req.should_close(): 193 self.req.unreader.close() 194 elif not self.closed: 195 # release connection 196 self.req.unreader.release() 197 self.closed = True
198
199 - def __enter__(self):
200 return self
201
202 - def __exit__(self, exc_type, exc_val, traceback):
203 if exc_type is None: 204 """ close on exit and release connection if needed """ 205 self.close()
206
207 - def __iter__(self):
208 return self
209
210 - def next(self):
211 ret = self.readline() 212 if not ret: 213 raise StopIteration() 214 return ret
215
216 - def getsize(self, size):
217 if size is None: 218 return sys.maxint 219 elif not isinstance(size, (int, long)): 220 raise TypeError("size must be an integral type") 221 elif size < 0: 222 return sys.maxint 223 return size
224
225 - def read(self, size=None):
226 size = self.getsize(size) 227 if size == 0: 228 return "" 229 230 if size < self.buf.tell(): 231 data = self.buf.getvalue() 232 ret, rest = data[:size], data[size:] 233 self.buf.truncate(0) 234 self.buf.write(rest) 235 return ret 236 237 while size > self.buf.tell(): 238 data = self.reader.read(1024) 239 if not len(data): 240 self.close() 241 break 242 self.buf.write(data) 243 244 data = self.buf.getvalue() 245 ret, rest = data[:size], data[size:] 246 self.buf.truncate(0) 247 self.buf.write(rest) 248 return ret
249
250 - def readline(self, size=None):
251 size = self.getsize(size) 252 if size == 0: 253 return "" 254 255 line = self.buf.getvalue() 256 idx = line.find("\n") 257 if idx >= 0: 258 ret = line[:idx+1] 259 self.buf.truncate(0) 260 self.buf.write(line[idx+1:]) 261 return ret 262 263 self.buf.truncate(0) 264 ch = "" 265 buf = [line] 266 lsize = len(line) 267 while lsize < size and ch != "\n": 268 ch = self.reader.read(1) 269 if not len(ch): 270 self.close() 271 break 272 lsize += 1 273 buf.append(ch) 274 return "".join(buf)
275
276 - def readlines(self, size=None):
277 ret = [] 278 data = self.read() 279 while len(data): 280 pos = data.find("\n") 281 if pos < 0: 282 ret.append(data) 283 data = "" 284 else: 285 line, data = data[:pos+1], data[pos+1:] 286 ret.append(line) 287 return ret
288 289
290 -class GzipBody(Body):
291 - def __init__(self, reader):
292 super(GzipBody, self).__init__(reader) 293 self._d = zlib.decompressobj(16+zlib.MAX_WBITS)
294
295 - def _decompress(self, data):
296 return self._d.decompress(data)
297
298 - def read(self, size=None):
299 size = self.getsize(size) 300 if size == 0: 301 return "" 302 303 if size < self.buf.tell(): 304 data = self.buf.getvalue() 305 ret, rest = data[:size], data[size:] 306 self.buf.truncate(0) 307 self.buf.write(rest) 308 return self._decompress(ret) 309 310 while size > self.buf.tell(): 311 data = self.reader.read(1024) 312 if not len(data): 313 self.close() 314 break 315 self.buf.write(data) 316 317 data = self.buf.getvalue() 318 ret, rest = data[:size], data[size:] 319 self.buf.truncate(0) 320 self.buf.write(rest) 321 return self._decompress(ret)
322
323 - def readline(self, size=None):
324 size = self.getsize(size) 325 if size == 0: 326 return "" 327 328 idx = self.buf.getvalue().find("\n") 329 while idx < 0: 330 data = self.reader.read(1024) 331 if not len(data): 332 self.close() 333 break 334 self.buf.write(self._decompress(data)) 335 idx = self.buf.getvalue().find("\n") 336 if size < self.buf.tell(): 337 break 338 339 # If we didn't find it, and we got here, we've 340 # exceeded size or run out of data. 341 if idx < 0: 342 rlen = min(size, self.buf.tell()) 343 else: 344 rlen = idx + 1 345 346 # If rlen is beyond our size threshold, trim back 347 if rlen > size: 348 rlen = size 349 350 data = self.buf.getvalue() 351 ret, rest = data[:rlen], data[rlen:] 352 353 self.buf.truncate(0) 354 self.buf.write(rest) 355 return ret
356 357
358 -class DeflateBody(GzipBody):
359 - def __init__(self, reader):
360 super(DeflateBody, self).__init__(reader) 361 self._d = zlib.decompressobj()
362