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 253 if not len(data): 254 self.close() 255 break 256 self.buf.write(data) 257 idx = self.buf.getvalue().find("\n") 258 if size < self.buf.tell(): 259 break 260 261 # If we didn't find it, and we got here, we've 262 # exceeded size or run out of data. 263 if idx < 0: 264 rlen = min(size, self.buf.tell()) 265 else: 266 rlen = idx + 1 267 268 # If rlen is beyond our size threshold, trim back 269 if rlen > size: 270 rlen = size 271 272 data = self.buf.getvalue() 273 ret, rest = data[:rlen], data[rlen:] 274 275 self.buf.truncate(0) 276 self.buf.write(rest) 277 return ret
278
279 - def readlines(self, size=None):
280 ret = [] 281 data = self.read() 282 while len(data): 283 pos = data.find("\n") 284 if pos < 0: 285 ret.append(data) 286 data = "" 287 else: 288 line, data = data[:pos+1], data[pos+1:] 289 ret.append(line) 290 return ret
291 292
293 -class GzipBody(Body):
294 - def __init__(self, reader):
295 Body.__init__(self, reader) 296 self._d = zlib.decompressobj(16+zlib.MAX_WBITS)
297
298 - def _decompress(self, data):
299 return self._d.decompress(data)
300
301 - def read(self, size=None):
302 size = self.getsize(size) 303 if size == 0: 304 return "" 305 306 if size < self.buf.tell(): 307 data = self.buf.getvalue() 308 ret, rest = data[:size], data[size:] 309 self.buf.truncate(0) 310 self.buf.write(rest) 311 return ret 312 313 while size > self.buf.tell(): 314 data = self.reader.read(1024) 315 if not len(data): 316 self.close() 317 break 318 data = self._decompress(data) 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 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 data = self._decompress(data) 339 self.buf.write(data) 340 idx = self.buf.getvalue().find("\n") 341 if size < self.buf.tell(): 342 break 343 344 # If we didn't find it, and we got here, we've 345 # exceeded size or run out of data. 346 if idx < 0: 347 rlen = min(size, self.buf.tell()) 348 else: 349 rlen = idx + 1 350 351 # If rlen is beyond our size threshold, trim back 352 if rlen > size: 353 rlen = size 354 355 data = self.buf.getvalue() 356 ret, rest = data[:rlen], data[rlen:] 357 358 self.buf.truncate(0) 359 self.buf.write(rest) 360 return ret
361