diff options
| -rw-r--r-- | wsgitools/middlewares.py | 66 |
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: |