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