refactor middlewares.NoWriteCallableMiddleware

1 files changed, 10 insertions, 7 deletions

diff --git a/wsgitools/middlewares.py b/wsgitools/middlewares.py
index 201634b..eb7b913 100644
--- a/wsgitools/middlewares.py
+++ b/wsgitools/middlewares.py
@@ -71,15 +71,20 @@ class NoWriteCallableMiddleware:
 
         ret = self.app(environ, modified_start_response)
         assert hasattr(ret, "__iter__")
+        assert len(todo) == 1
 
-        if todo and todo[0] is None:
+        if todo[0] is None:
+            return ret
+
+        status, headers, data = todo[0]
+        data = data.getvalue()
+
+        if not data:
+            start_response(status, headers)
             return ret
 
         if isinstance(ret, list):
-            status, headers, data = todo[0]
-            data = data.getvalue()
-            if data:
-                ret.insert(0, data)
+            ret.insert(0, data)
             start_response(status, headers)
             return ret
 
@@ -90,8 +95,6 @@ class NoWriteCallableMiddleware:
         except StopIteration:
             stopped = True
 
-        status, headers, data = todo[0]
-        data = data.getvalue()
         start_response(status, headers)
 
         if stopped: