1 files changed, 36 insertions, 30 deletions
diff --git a/wsgitools/middlewares.py b/wsgitools/middlewares.py
index b385e91..fa2de2e 100644
--- a/wsgitools/middlewares.py
+++ b/wsgitools/middlewares.py
@@ -69,51 +69,57 @@ class NoWriteCallableMiddleware:
         @rtype: gen([str])
         """
         assert isinstance(environ, dict)
-        todo = []
+        todo = [None]
+        sio = io.StringIO()
+        gotiterdata = False
+        def write_calleable(data):
+            assert not gotiterdata
+            sio.write(data)
         def modified_start_response(status, headers, exc_info=None):
+            try:
+                if sio.tell() > 0 or gotiterdata:
+                    assert exc_info is not None
+                    raise exc_info[0], exc_info[1], exc_info[2]
+            finally:
+                exc_info = None
             assert isinstance(status, str)
             assert isinstance(headers, list)
-            if exc_info is not None:
-                todo.append(None)
-                return start_response(status, headers, exc_info)
-            else:
-                sio = io.StringIO()
-                todo.append((status, headers, sio))
-                return sio.write
+            todo[0] = (status, headers)
+            return write_calleable
 
         ret = self.app(environ, modified_start_response)
         assert hasattr(ret, "__iter__")
-        assert len(todo) == 1
 
-        if todo[0] is None:
-            return ret
-
-        status, headers, data = todo[0]
-        data = data.getvalue()
-
-        if not data:
-            start_response(status, headers)
-            return ret
+        first = ""
+        if not isinstance(ret, list):
+            ret = iter(ret)
+            stopped = False
+            while not (stopped or first):
+                try:
+                    first = next(ret)
+                except StopIteration:
+                    stopped = True
+            gotiterdata = True
+            if stopped:
+                ret = CloseableList(getattr(ret, "close", None), (first,))
+        else:
+            gotiterdata = True
+
+        assert todo[0] is not None
+        status, headers = todo[0]
+        data = sio.getvalue()
 
         if isinstance(ret, list):
-            ret.insert(0, data)
+            if data:
+                ret.insert(0, data)
             start_response(status, headers)
             return ret
 
-        ret = iter(ret)
-        stopped = False
-        try:
-            first = next(ret)
-        except StopIteration:
-            stopped = True
-
+        data += first
         start_response(status, headers)
 
-        if stopped:
-            return CloseableList(getattr(ret, "close", None), (data,))
-
         return CloseableIterator(getattr(ret, "close", None),
-                                 (data, first), ret)
+                                 (data,), ret)
 
 __all__.append("ContentLengthMiddleware")
 class ContentLengthMiddleware: