Support for closing filelike response stream objects

falcon/api.py

@@ -680,10 +680,7 @@ def _get_body(self, resp, wsgi_file_wrapper=None):
                     iterable = wsgi_file_wrapper(stream,
                                                  self._STREAM_BLOCK_SIZE)
                 else:
-                    iterable = iter(
-                        lambda: stream.read(self._STREAM_BLOCK_SIZE),
-                        b''
-                    )
+                    iterable = helpers.CloseableStreamIterator(stream, self._STREAM_BLOCK_SIZE)
             else:
                 iterable = stream
 

falcon/api_helpers.py

@@ -16,6 +16,8 @@
 
 from functools import wraps
 
+import six
+
 from falcon import util
 
 
@@ -167,3 +169,40 @@ def new_fn(req, resp, exception):
             resp.content_type = media_type
 
     return new_fn
+
+
+class CloseableStreamIterator(six.Iterator):
+    """
+    Iterator that wraps stream and supports closing underlying stream.
+
+    This iterator can be used to read from an underlying stream in
+    block_size-chunks until response from stream is empty string (bytes).
+
+    This class is used to wrap wsgi response streams (when no explicit
+    wsgi_file_wrapper is specified).  The fact that it also supports closing
+    the underlying stream allows use of (e.g.) python tempfile resources that
+    would be deleted upon close.
+    """
+
+    def __init__(self, stream, block_size):
+        """
+        Args:
+            stream: Stream file-like object.
+            block_size: Number of bytes to read per iteration.
+        """
+        self.stream = stream
+        self.block_size = block_size
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        data = self.stream.read(self.block_size)
+        if data == b'':
+            raise StopIteration
+        else:
+            return data
+
+    def close(self):
+        if hasattr(self.stream, 'close') and callable(self.stream.close):
+            self.stream.close()

tests/test_hello.py

@@ -69,6 +69,34 @@ def on_head(self, req, resp):
         self.on_get(req, resp)
 
 
+class ClosingBytesIO(io.BytesIO):
+
+    close_called = False
+
+    def close(self):
+        super(ClosingBytesIO, self).close()
+        self.close_called = True
+
+
+class ClosingFilelikeHelloResource(object):
+    sample_status = '200 OK'
+    sample_unicode = (u'Hello World! \x80' +
+                      six.text_type(testing.rand_string(0, 0)))
+
+    sample_utf8 = sample_unicode.encode('utf-8')
+
+    def __init__(self):
+        self.called = False
+        self.stream = ClosingBytesIO(self.sample_utf8)
+        self.stream_len = len(self.sample_utf8)
+
+    def on_get(self, req, resp):
+        self.called = True
+        self.req, self.resp = req, resp
+        resp.status = falcon.HTTP_200
+        resp.set_stream(self.stream, self.stream_len)
+
+
 class NoStatusResource(object):
     def on_get(self, req, resp):
         pass
@@ -158,6 +186,28 @@ def test_filelike(self):
             self.assertEqual(actual_len, expected_len)
             self.assertEqual(len(result.content), expected_len)
 
+        for file_wrapper in (None, FileWrapper):
+            result = self.simulate_get('/filelike', file_wrapper=file_wrapper)
+            self.assertTrue(resource.called)
+
+            expected_len = resource.resp.stream_len
+            actual_len = int(result.headers['content-length'])
+            self.assertEqual(actual_len, expected_len)
+            self.assertEqual(len(result.content), expected_len)
+
+    def test_filelike_closing(self):
+        resource = ClosingFilelikeHelloResource()
+        self.api.add_route('/filelike-closing', resource)
+
+        result = self.simulate_get('/filelike-closing', file_wrapper=None)
+        self.assertTrue(resource.called)
+
+        expected_len = resource.resp.stream_len
+        actual_len = int(result.headers['content-length'])
+        self.assertEqual(actual_len, expected_len)
+        self.assertEqual(len(result.content), expected_len)
+        self.assertTrue(resource.stream.close_called)
+
     def test_filelike_using_helper(self):
         resource = HelloResource('stream, stream_len, filelike, use_helper')
         self.api.add_route('/filelike-helper', resource)