From 08389c92a91c3018c46086fb4037a398303256f4 Mon Sep 17 00:00:00 2001
From: Hans Lellelid <hans@xmpl.org>
Date: Fri, 27 Jan 2017 15:23:06 -0500
Subject: [PATCH] feat(API): Support for closing filelike response stream
 objects (#961)

Fixes #960
---
 falcon/api.py         |  5 +---
 falcon/api_helpers.py | 38 ++++++++++++++++++++++++++
 tests/test_hello.py   | 63 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 4 deletions(-)

diff --git a/falcon/api.py b/falcon/api.py
index 659eeec48..1e057f222 100644
--- a/falcon/api.py
+++ b/falcon/api.py
@@ -690,10 +690,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
 
diff --git a/falcon/api_helpers.py b/falcon/api_helpers.py
index e799c9746..550f87329 100644
--- a/falcon/api_helpers.py
+++ b/falcon/api_helpers.py
@@ -16,6 +16,8 @@
 
 from functools import wraps
 
+import six
+
 from falcon import util
 
 
@@ -167,3 +169,39 @@ def new_fn(req, resp, exception):
             resp.content_type = media_type
 
     return new_fn
+
+
+class CloseableStreamIterator(six.Iterator):
+    """Iterator that wraps a file-like stream with support for close().
+
+    This iterator can be used to read from an underlying file-like stream
+    in block_size-chunks until the response from the stream is an empty
+    byte string.
+
+    This class is used to wrap WSGI response streams when a
+    wsgi_file_wrapper is not provided by the server.  The fact that it
+    also supports closing the underlying stream allows use of (e.g.)
+    Python tempfile resources that would be deleted upon close.
+
+    Args:
+        stream (object): Readable file-like stream object.
+        block_size (int): Number of bytes to read per iteration.
+    """
+
+    def __init__(self, stream, block_size):
+        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()
diff --git a/tests/test_hello.py b/tests/test_hello.py
index ceedb1a6d..be41ddff9 100644
--- a/tests/test_hello.py
+++ b/tests/test_hello.py
@@ -69,6 +69,40 @@ 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 NonClosingBytesIO(io.BytesIO):
+
+    # Not callable; test that CloseableStreamIterator ignores it
+    close = False
+
+
+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, stream_factory):
+        self.called = False
+        self.stream = stream_factory(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 +192,35 @@ 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)
+
+    @ddt.data(
+        (ClosingBytesIO, True),  # Implements close()
+        (NonClosingBytesIO, False),  # Has a non-callable "close" attr
+    )
+    @ddt.unpack
+    def test_filelike_closing(self, stream_factory, assert_closed):
+        resource = ClosingFilelikeHelloResource(stream_factory)
+        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)
+
+        if assert_closed:
+            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)