defnull avatar defnull committed 178428d

Iterators with a .close() method are now wrapped in a way that preserves the .close() method.

Comments (0)

Files changed (2)

 
         # Handle Iterables. We peek into them to detect their inner type.
         try:
-            out = iter(out)
-            first = next(out)
+            iout = iter(out)
+            first = next(iout)
             while not first:
-                first = next(out)
+                first = next(iout)
         except StopIteration:
             return self._cast('')
         except HTTPResponse:
         # These are the inner types allowed in iterator or generator objects.
         if isinstance(first, HTTPResponse):
             return self._cast(first)
-        if isinstance(first, bytes):
-            return itertools.chain([first], out)
-        if isinstance(first, unicode):
-            return imap(lambda x: x.encode(response.charset),
-                                  itertools.chain([first], out))
-        return self._cast(HTTPError(500, 'Unsupported response type: %s'\
-                                         % type(first)))
+        elif isinstance(first, bytes):
+            new_iter = itertools.chain([first], iout)
+        elif isinstance(first, unicode):
+            encoder = lambda x: x.encode(response.charset)
+            new_iter = imap(encoder, itertools.chain([first], iout))
+        else:
+            msg = 'Unsupported response type: %s' % type(first)
+            return self._cast(HTTPError(500, msg))
+        if hasattr(out, 'close'):
+            new_iter = _iterchain(new_iter)
+            new_iter.close = out.close
+        return new_iter
 
     def wsgi(self, environ, start_response):
         """ The bottle WSGI-interface. """
             yield part
 
 
+class _iterchain(itertools.chain):
+    ''' This only exists to be able to attach a .close method to iterators that
+        do not support attribute assignment (most of itertools). '''
+
+
 class ResourceManager(object):
     ''' This class manages a list of search paths and helps to find and open
         application-bound resources (files).

test/test_outputfilter.py

         self.assertStatus(500)
         self.assertInBody('Unsupported response type')
 
+    def test_iterator_with_close(self):
+        class MyIter(object):
+            def __init__(self, data):
+                self.data = data
+                self.closed = False
+            def close(self):    self.closed = True
+            def __iter__(self): return iter(self.data)
+
+        byte_iter = MyIter([tob('abc'), tob('def')])
+        unicode_iter = MyIter([touni('abc'), touni('def')])
+
+        for test_iter in (byte_iter, unicode_iter):
+            @self.app.route('/')
+            def test(): return byte_iter
+            self.assertInBody('abcdef')
+            self.assertTrue(byte_iter.closed)
+
     def test_cookie(self):
         """ WSGI: Cookies """
         @bottle.route('/cookie')
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.