summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--wsgitools/middlewares.py66
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: