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 or self.finished: 158 return "" 159 160 161 data = self.unreader.read() 162 while data: 163 self.buf.write(data) 164 if self.buf.tell() > size: 165 break 166 data = self.unreader.read() 167 168 if not data: 169 self.finished = True 170 return self.buf.getvalue() 171 172 data = self.buf.getvalue() 173 ret, rest = data[:size], data[size:] 174 self.buf.truncate(0) 175 self.buf.write(rest) 176 return ret
177
178 -class Body(object):
179 - def __init__(self, reader):
180 self.reader = reader 181 self.req = reader.req 182 self.buf = StringIO() 183 self.closed = False
184
185 - def close(self):
186 """ Close the socket if needed """ 187 if self.req.should_close(): 188 self.req.unreader.close() 189 elif not self.closed: 190 # release connection 191 self.req.unreader.release() 192 self.closed = True
193
194 - def __enter__(self):
195 return self
196
197 - def __exit__(self, exc_type, exc_val, traceback):
198 if exc_type is None: 199 """ close on exit and release connection if needed """ 200 self.close()
201
202 - def __iter__(self):
203 return self
204
205 - def next(self):
206 ret = self.readline() 207 if not ret: 208 raise StopIteration() 209 return ret
210
211 - def getsize(self, size):
212 if size is None: 213 return sys.maxint 214 elif not isinstance(size, (int, long)): 215 raise TypeError("size must be an integral type") 216 elif size < 0: 217 return sys.maxint 218 return size
219
220 - def read(self, size=None):
221 size = self.getsize(size) 222 if size == 0: 223 return "" 224 225 if size < self.buf.tell(): 226 data = self.buf.getvalue() 227 ret, rest = data[:size], data[size:] 228 self.buf.truncate(0) 229 self.buf.write(rest) 230 return ret 231 232 while size > self.buf.tell(): 233 data = self.reader.read(1024) 234 if not len(data): 235 self.close() 236 break 237 self.buf.write(data) 238 239 data = self.buf.getvalue() 240 ret, rest = data[:size], data[size:] 241 self.buf.truncate(0) 242 self.buf.write(rest) 243 return ret
244
245 - def readline(self, size=None):
246 size = self.getsize(size) 247 if size == 0: 248 return "" 249 250 idx = self.buf.getvalue().find("\n") 251 while idx < 0: 252 data = self.reader.read(1024) 253 254 if not len(data): 255 self.close() 256 break 257 self.buf.write(data) 258 idx = self.buf.getvalue().find("\n") 259 if size < self.buf.tell(): 260 break 261 262 # If we didn't find it, and we got here, we've 263 # exceeded size or run out of data. 264 if idx < 0: 265 rlen = min(size, self.buf.tell()) 266 else: 267 rlen = idx + 1 268 269 # If rlen is beyond our size threshold, trim back 270 if rlen > size: 271 rlen = size 272 273 data = self.buf.getvalue() 274 ret, rest = data[:rlen], data[rlen:] 275 276 self.buf.truncate(0) 277 self.buf.write(rest) 278 return ret
279
280 - def readlines(self, size=None):
281 ret = [] 282 data = self.read() 283 while len(data): 284 pos = data.find("\n") 285 if pos < 0: 286 ret.append(data) 287 data = "" 288 else: 289 line, data = data[:pos+1], data[pos+1:] 290 ret.append(line) 291 return ret
292 293
294 -class GzipBody(Body):
295 - def __init__(self, reader):
296 super(GzipBody, self).__init__(reader) 297 self._d = zlib.decompressobj(16+zlib.MAX_WBITS)
298
299 - def _decompress(self, data):
300 return self._d.decompress(data)
301
302 - def read(self, size=None):
303 size = self.getsize(size) 304 if size == 0: 305 return "" 306 307 if size < self.buf.tell(): 308 data = self.buf.getvalue() 309 ret, rest = data[:size], data[size:] 310 self.buf.truncate(0) 311 self.buf.write(rest) 312 return self._decompress(ret) 313 314 while size > self.buf.tell(): 315 data = self.reader.read(1024) 316 if not len(data): 317 self.close() 318 break 319 self.buf.write(data) 320 321 data = self.buf.getvalue() 322 ret, rest = data[:size], data[size:] 323 self.buf.truncate(0) 324 self.buf.write(rest) 325 return self._decompress(ret)
326
327 - def readline(self, size=None):
328 size = self.getsize(size) 329 if size == 0: 330 return "" 331 332 idx = self.buf.getvalue().find("\n") 333 while idx < 0: 334 data = self.reader.read(1024) 335 if not len(data): 336 self.close() 337 break 338 self.buf.write(self._decompress(data)) 339 idx = self.buf.getvalue().find("\n") 340 if size < self.buf.tell(): 341 break 342 343 # If we didn't find it, and we got here, we've 344 # exceeded size or run out of data. 345 if idx < 0: 346 rlen = min(size, self.buf.tell()) 347 else: 348 rlen = idx + 1 349 350 # If rlen is beyond our size threshold, trim back 351 if rlen > size: 352 rlen = size 353 354 data = self.buf.getvalue() 355 ret, rest = data[:rlen], data[rlen:] 356 357 self.buf.truncate(0) 358 self.buf.write(rest) 359 return ret
360 361
362 -class DeflateBody(GzipBody):
363 - def __init__(self, reader):
364 super(DeflateBody, self).__init__(reader) 365 self._d = zlib.decompressobj()
366