Simplified iteration in HTTP response objects.

Fixed #20187 -- Allowed repeated iteration of HttpResponse.

All this became possible when support for old-style streaming responses was
finally removed.
This commit is contained in:
Aymeric Augustin 2013-11-23 17:03:43 +01:00
parent 1124e16340
commit a480f8320a
2 changed files with 14 additions and 16 deletions

View File

@ -282,13 +282,6 @@ class HttpResponseBase(six.Iterator):
# Handle non-string types (#16494) # Handle non-string types (#16494)
return force_bytes(value, self._charset) return force_bytes(value, self._charset)
def __iter__(self):
return self
def __next__(self):
# Subclasses must define self._iterator for this function.
return self.make_bytes(next(self._iterator))
# These methods partially implement the file-like object interface. # These methods partially implement the file-like object interface.
# See http://docs.python.org/lib/bltin-file-objects.html # See http://docs.python.org/lib/bltin-file-objects.html
@ -337,23 +330,25 @@ class HttpResponse(HttpResponseBase):
@property @property
def content(self): def content(self):
return b''.join(self.make_bytes(e) for e in self._container) return b''.join(self._container)
@content.setter @content.setter
def content(self, value): def content(self, value):
# Consume iterators upon assignment to allow repeated iteration.
if hasattr(value, '__iter__') and not isinstance(value, (bytes, six.string_types)): if hasattr(value, '__iter__') and not isinstance(value, (bytes, six.string_types)):
if hasattr(value, 'close'): if hasattr(value, 'close'):
self._closable_objects.append(value) self._closable_objects.append(value)
value = b''.join(self.make_bytes(e) for e in value) value = b''.join(self.make_bytes(chunk) for chunk in value)
else:
value = self.make_bytes(value)
# Create a list of properly encoded bytestrings to support write().
self._container = [value] self._container = [value]
def __iter__(self): def __iter__(self):
if not hasattr(self, '_iterator'): return iter(self._container)
self._iterator = iter(self._container)
return self
def write(self, content): def write(self, content):
self._container.append(content) self._container.append(self.make_bytes(content))
def tell(self): def tell(self):
return len(self.content) return len(self.content)
@ -392,6 +387,9 @@ class StreamingHttpResponse(HttpResponseBase):
if hasattr(value, 'close'): if hasattr(value, 'close'):
self._closable_objects.append(value) self._closable_objects.append(value)
def __iter__(self):
return self.streaming_content
class HttpResponseRedirectBase(HttpResponse): class HttpResponseRedirectBase(HttpResponse):
allowed_schemes = ['http', 'https', 'ftp'] allowed_schemes = ['http', 'https', 'ftp']

View File

@ -356,10 +356,10 @@ class HttpResponseTests(unittest.TestCase):
r = HttpResponse(iter(['hello', 'world'])) r = HttpResponse(iter(['hello', 'world']))
self.assertEqual(r.content, r.content) self.assertEqual(r.content, r.content)
self.assertEqual(r.content, b'helloworld') self.assertEqual(r.content, b'helloworld')
# accessing the iterator works (once) after accessing .content # __iter__ can safely be called multiple times (#20187).
self.assertEqual(b''.join(r), b'helloworld') self.assertEqual(b''.join(r), b'helloworld')
self.assertEqual(b''.join(r), b'') self.assertEqual(b''.join(r), b'helloworld')
# accessing .content still works # Accessing .content still works.
self.assertEqual(r.content, b'helloworld') self.assertEqual(r.content, b'helloworld')
# Accessing .content also works if the response was iterated first. # Accessing .content also works if the response was iterated first.