Close iterable after iteration

.travis.yml

@@ -12,7 +12,7 @@ env:
   - TOXENV=py33
   - TOXENV=py34
   - TOXENV=pypy
-  - TOXENV=pypy3
+  - TOXENV=pypy3.3-5.2-alpha1
   - TOXENV=cover
 
 install:

CONTRIBUTORS.txt

@@ -106,3 +106,4 @@ Contributors
 - Tres Seaver, 2011/02/22
 - Stéphane Klein <[email protected]>, 2011/06/18
 - Brian Sutherland, 2011/07/11
+- Mike Milner <[email protected]>, 2017-03-21

repoze/tm/__init__.py

@@ -36,7 +36,7 @@ def __init__(self, application, commit_veto=None):
         self.application = application
         self.commit_veto = commit_veto
         self.transaction = transaction # for testing
-        
+
     def __call__(self, environ, start_response):
         transaction = self.transaction
         environ[ekey] = True
@@ -48,8 +48,16 @@ def save_status_and_headers(status, headers, exc_info=None):
             return start_response(status, headers, exc_info)
 
         try:
-            for chunk in self.application(environ, save_status_and_headers):
-                yield chunk
+            response = self.application(environ, save_status_and_headers)
+            # Ensure that the `.close()` method is always called (if present)
+            # after iteration. Required by WSGI specification.
+            # https://www.python.org/dev/peps/pep-0333/#specification-details
+            try:
+                for chunk in response:
+                    yield chunk
+            finally:
+                if hasattr(response, "close"):
+                    response.close()
         except Exception:
             """Saving the exception"""
             try:

repoze/tm/tests.py

@@ -52,7 +52,7 @@ def execute_request():
         self.assertRaises(ValueError, execute_request)
         self.assertEqual(transaction.committed, False)
         self.assertEqual(transaction.aborted, True)
-        
+
     def test_aborted_via_exception_and_doom(self):
         app = DummyApplication(exception=True)
         tm = self._makeOne(app)
@@ -122,7 +122,7 @@ def dummy():
         self.assertEqual(transaction.committed, True)
         self.assertEqual(transaction.aborted, False)
         self.assertEqual(dummycalled, [True])
-        
+
     def test_cleanup_on_abort(self):
         from repoze.tm import after_end
         dummycalled = []
@@ -142,6 +142,56 @@ def execute_request():
         self.assertEqual(transaction.aborted, True)
         self.assertEqual(dummycalled, [True])
 
+    def test_response_iter_closed(self):
+        """
+        If the app returns an iterator with a `.close()` method, it is
+        called after consuming the iterator.
+        """
+        class DummyResponse(object):
+            def __init__(self):
+                self.closed = False
+
+            def __iter__(self):
+                yield "hello"
+                yield "iter"
+
+            def close(self):
+                self.closed = True
+
+        env = {}
+        response = DummyResponse()
+        app = DummyApplication(response=response)
+        tm = self._makeOne(app)
+        result = [chunk for chunk in tm({}, self._start_response)]
+        self.assertEqual(result, ["hello", "iter"])
+        self.assertTrue(response.closed)
+
+    def test_response_iter_closed_on_error(self):
+        """
+        If the app returns an iterator with a `.close()` method, it is
+        called even if the iteration raises an error.
+        """
+        class DummyResponse(object):
+            def __init__(self):
+                self.closed = False
+
+            def __iter__(self):
+                yield "hello"
+                raise Exception("BOOM!")
+
+            def close(self):
+                self.closed = True
+
+        env = {}
+        response = DummyResponse()
+        app = DummyApplication(response=response)
+        tm = self._makeOne(app)
+        def execute_request():
+            [chunk for chunk in tm({}, self._start_response)]
+        self.assertRaises(Exception, execute_request)
+        self.assertTrue(response.closed)
+
+
 class TestAfterEnd(unittest.TestCase):
     def _getTargetClass(self):
         from repoze.tm import AfterEnd
@@ -165,7 +215,7 @@ def test_unregister_exists(self):
         self.assertEqual(getattr(txn, registry.key), [func])
         registry.unregister(func, txn)
         self.assertFalse(hasattr(txn, registry.key))
-        
+
     def test_unregister_notexists(self):
         registry = self._makeOne()
         func = lambda *x: None
@@ -202,7 +252,7 @@ class Test_default_commit_veto(unittest.TestCase):
     def _callFUT(self, status, headers=()):
         from repoze.tm import default_commit_veto
         return default_commit_veto(None, status, headers)
-    
+
     def test_it_true_5XX(self):
         self.assertTrue(self._callFUT('500 Server Error'))
         self.assertTrue(self._callFUT('503 Service Unavailable'))
@@ -264,13 +314,16 @@ class Dummy:
     pass
 
 class DummyApplication:
-    def __init__(self, exception=False, status="200 OK"):
+    def __init__(self, exception=False, status="200 OK", response=None):
         self.exception = exception
         self.status = status
-
+        self.response = response
+
     def __call__(self, environ, start_response):
         start_response(self.status, [], None)
         if self.exception:
             raise ValueError('raising')
-        return ['hello']
-
+        if self.response is None:
+            return ['hello']
+        else:
+            return self.response

tox.ini

@@ -1,6 +1,6 @@
 [tox]
 envlist = 
-    py27,pypy,py33,py34,py35,pypy3,cover,docs
+    py27,pypy,py33,py34,py35,pypy3.3-5.2-alpha1,cover,docs
 
 [testenv]
 commands =