Merge pull request #150
authorKenneth Reitz <me@kennethreitz.com>
Wed, 14 Sep 2011 12:39:52 +0000 (08:39 -0400)
committerKenneth Reitz <me@kennethreitz.com>
Wed, 14 Sep 2011 12:39:52 +0000 (08:39 -0400)
requests/models.py
requests/utils.py
test_requests.py

index 6180e969d8a6c873fab7c20a1f17c260365df3e3..a16d32eae3fc82429b2e702d895546918ad01565 100644 (file)
@@ -9,6 +9,7 @@ requests.models
 import urllib
 import urllib2
 import socket
+import codecs
 import zlib
 
 
@@ -21,7 +22,7 @@ from .monkeys import Request as _Request, HTTPBasicAuthHandler, HTTPForcedBasicA
 from .structures import CaseInsensitiveDict
 from .packages.poster.encode import multipart_encode
 from .packages.poster.streaminghttp import register_openers, get_handlers
-from .utils import dict_from_cookiejar, get_unicode_from_response, decode_gzip
+from .utils import dict_from_cookiejar, get_unicode_from_response, stream_decode_response_unicode, decode_gzip, stream_decode_gzip
 from .status_codes import codes
 from .exceptions import RequestException, AuthenticationError, Timeout, URLRequired, InvalidMethod, TooManyRedirects
 
@@ -396,6 +397,7 @@ class Response(object):
     def __init__(self):
 
         self._content = None
+        self._content_consumed = False
 
         #: Integer Code of responded HTTP Status.
         self.status_code = None
@@ -438,6 +440,31 @@ class Response(object):
 
         return not self.error
 
+    def iter_content(self, chunk_size=10 * 1024, decode_unicode=None):
+        """Iterates over the response data.  This avoids reading the content
+        at once into memory for large responses.  The chunk size is the number
+        of bytes it should read into memory.  This is not necessarily the
+        length of each item returned as decoding can take place.
+        """
+        if self._content_consumed:
+            raise RuntimeError('The content for this response was '
+                               'already consumed')
+
+        def generate():
+            while 1:
+                chunk = self.fo.read(chunk_size)
+                if not chunk:
+                    break
+                yield chunk
+            self._content_consumed = True
+        gen = generate()
+        if 'gzip' in self.headers.get('content-encoding', ''):
+            gen = stream_decode_gzip(gen)
+        if decode_unicode is None:
+            decode_unicode = settings.decode_unicode
+        if decode_unicode:
+            gen = stream_decode_response_unicode(gen, self)
+        return gen
 
     @property
     def content(self):
@@ -448,6 +475,10 @@ class Response(object):
         if self._content is not None:
             return self._content
 
+        if self._content_consumed:
+            raise RuntimeError('The content for this response was '
+                               'already consumed')
+
         # Read the contents.
         self._content = self.fo.read()
 
@@ -462,6 +493,7 @@ class Response(object):
         if settings.decode_unicode:
             self._content = get_unicode_from_response(self)
 
+        self._content_consumed = True
         return self._content
 
 
index c1aed434fc3dfbc94b01e84f8639c9ea77ab62bf..17f79a8ac8e2ffc423985bcd52effa2f3d7d6221 100644 (file)
@@ -10,6 +10,7 @@ that are also useful for external consumption.
 """
 
 import cgi
+import codecs
 import cookielib
 import re
 import zlib
@@ -177,6 +178,24 @@ def unicode_from_html(content):
         return content
 
 
+def stream_decode_response_unicode(iterator, r):
+    """Stream decodes a iterator."""
+    encoding = get_encoding_from_headers(r.headers)
+    if encoding is None:
+        for item in iterator:
+            yield item
+        return
+
+    decoder = codecs.getincrementaldecoder(encoding)(errors='replace')
+    for chunk in iterator:
+        rv = decoder.decode(chunk)
+        if rv:
+            yield rv
+    rv = decoder.decode('', final=True)
+    if rv:
+        yield rv
+
+
 def get_unicode_from_response(r):
     """Returns the requested content back in unicode.
 
@@ -217,7 +236,24 @@ def decode_gzip(content):
     """
 
     return zlib.decompress(content, 16+zlib.MAX_WBITS)
+    return zlib.decompress(content, 16+zlib.MAX_WBITS)
+    return zlib.decompress(content, 16 + zlib.MAX_WBITS)
+
 
+def stream_decode_gzip(iterator):
+    """Stream decodes a gzip-encoded iterator"""
+    try:
+        dec = zlib.decompressobj(16 + zlib.MAX_WBITS)
+        for chunk in iterator:
+            rv = dec.decompress(chunk)
+            if rv:
+                yield rv
+        buf = dec.decompress('')
+        rv = buf + dec.flush()
+        if rv:
+            yield rv
+    except zlib.error:
+        pass
 
 def curl_from_request(request):
     """Returns a curl command from the request.
index 5baff568262d269c27fd20f28132442141a2aa25..406d23c19ab8e8675c47c8c55c5da5385e6ce779 100755 (executable)
@@ -478,6 +478,7 @@ class RequestsTestSuite(unittest.TestCase):
         self.assertEqual(r2.status_code, 200)
 
 
+
     def test_curl_HTTP_OK_GET(self):
         curl_str = 'curl -L -X GET -H "Accept-Encoding:gzip" -H "User-Agent:python-requests.org" "http://httpbin.org//"'
         r = requests.get(httpbin('/'))