Merge pull request #28 from gweis/master

support IStreamIterator as WSGI response body

CHANGES.rst

@@ -29,6 +29,12 @@ Bugs Fixed
 
 - Made sure getConfiguration().default_zpublisher_encoding is set correctly.
 
+- Issue #28: Fix publishing of IStreamIterator. This interface does
+  not have seek or tell.
+  Introduce IUnboundStreamIterator to support publishing iterators
+  of unknown length.
+
+
 Features Added
 ++++++++++++++
 

src/ZPublisher/Iterators.py

 from zope.interface import Interface
 from zope.interface import implements
 
-class IStreamIterator(Interface):
+class IUnboundStreamIterator(Interface):
     """
-    An iterator that can be published.
+    An iterator with unknown length that can be published.
+    """
+
+    def next():
+        """
+        Return a sequence of bytes out of the bytestream, or raise
+        StopIeration if we've reached the end of the bytestream.
+        """
+
+
+class IStreamIterator(IUnboundStreamIterator):
+    """
+    An iterator with known length that can be published.
 
     IStreamIterators must not read from the object database.
     After the application finishes interpreting a request and
@@ -15,12 +27,6 @@ class IStreamIterator(Interface):
     but it has a chance of going insane if it happens to be loading
     or storing something in the other thread at the same time.                      """
 
-    def next():
-        """
-        Return a sequence of bytes out of the bytestream, or raise
-        StopIeration if we've reached the end of the bytestream.
-        """
-
     def __len__():
         """
         Return an integer representing the length of the object

src/ZPublisher/WSGIPublisher.py

@@ -32,7 +32,7 @@ from ZPublisher.Publish import dont_publish_class
 from ZPublisher.Publish import get_module_info
 from ZPublisher.Publish import missing_name
 from ZPublisher.pubevents import PubStart, PubBeforeCommit, PubAfterTraversal
-from ZPublisher.Iterators import IStreamIterator
+from ZPublisher.Iterators import IUnboundStreamIterator, IStreamIterator
 
 _NOW = None     # overwrite for testing
 def _now():
@@ -131,12 +131,19 @@ class WSGIResponse(HTTPResponse):
         self.stdout.write(data)
 
     def setBody(self, body, title='', is_error=0):
-        if isinstance(body, file) or IStreamIterator.providedBy(body):
+        if isinstance(body, file):
             body.seek(0, 2)
             length = body.tell()
             body.seek(0)
             self.setHeader('Content-Length', '%d' % length)
             self.body = body
+        elif IStreamIterator.providedBy(body):
+            self.body = body
+            HTTPResponse.setBody(self, '', title, is_error)
+        elif IUnboundStreamIterator.providedBy(body):
+            self.body = body
+            self._streaming = 1
+            HTTPResponse.setBody(self, '', title, is_error)
         else:
             HTTPResponse.setBody(self, body, title, is_error)
 

src/ZPublisher/tests/test_WSGIPublisher.py

@@ -136,6 +136,57 @@ class WSGIResponseTests(unittest.TestCase):
                                 time.gmtime(time.mktime(WHEN)))
         self.assertTrue(('Date', whenstr) in headers)
 
+    def test_setBody_IUnboundStreamIterator(self):
+        from ZPublisher.Iterators import IUnboundStreamIterator
+        from zope.interface import implements
+
+        class test_streamiterator:
+            implements(IUnboundStreamIterator)
+            data = "hello"
+            done = 0
+
+            def next(self):
+                if not self.done:
+                    self.done = 1
+                    return self.data
+                raise StopIteration
+
+        response = self._makeOne()
+        response.setStatus(200)
+        body = test_streamiterator()
+        response.setBody(body)
+        response.finalize()
+        self.assertTrue(body is response.body)
+        self.assertEqual(response._streaming, 1)
+
+    def test_setBody_IStreamIterator(self):
+        from ZPublisher.Iterators import IStreamIterator
+        from zope.interface import implements
+
+        class test_streamiterator:
+            implements(IStreamIterator)
+            data = "hello"
+            done = 0
+
+            def next(self):
+                if not self.done:
+                    self.done = 1
+                    return self.data
+                raise StopIteration
+
+            def __len__(self):
+                return len(self.data)
+
+        response = self._makeOne()
+        response.setStatus(200)
+        body = test_streamiterator()
+        response.setBody(body)
+        response.finalize()
+        self.assertTrue(body is response.body)
+        self.assertEqual(response._streaming, 0)
+        self.assertEqual(response.getHeader('Content-Length'),
+                         '%d' % len(test_streamiterator.data))
+
     #def test___str__already_wrote_not_chunking(self):
     #    response = self._makeOne()
     #    response._wrote = True