TracebackMiddleware catches more exceptions

1 files changed, 11 insertions, 1 deletions

diff --git a/wsgitools/middlewares.py b/wsgitools/middlewares.py
index ffb2e38..a0b4a7c 100644
--- a/wsgitools/middlewares.py
+++ b/wsgitools/middlewares.py
@@ -278,7 +278,17 @@ class TracebackMiddleware:
     def __call__(self, environ, start_response):
         """wsgi interface"""
         try:
-            return self.app(environ, start_response)
+            ret = self.app(environ, start_response)
+            if isinstance(ret, list):
+                return ret
+            # Take the first element of the iterator and possibly catch an
+            # exception there.
+            ret = iter(ret)
+            try:
+                first = ret.next()
+            except StopIteration:
+                return CloseableList(getattr(ret, "close", None), [])
+            return CloseableIterator(getattr(ret, "close", None), [first], ret)
         except:
             exc_info = sys.exc_info()
             data = cgitb.html(exc_info)