diff options
| -rw-r--r-- | wsgitools/middlewares.py | 62 |
1 files changed, 34 insertions, 28 deletions
diff --git a/wsgitools/middlewares.py b/wsgitools/middlewares.py index fa2de2e..9792089 100644 --- a/wsgitools/middlewares.py +++ b/wsgitools/middlewares.py @@ -143,45 +143,49 @@ class ContentLengthMiddleware: """wsgi interface""" assert isinstance(environ, dict) todo = [] + gotdata = False def modified_start_response(status, headers, exc_info=None): + try: + if gotdata: + 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 or - [v for h, v in headers if h.lower() == "content-length"]): - todo[:] = (None,) - return start_response(status, headers, exc_info) - else: - todo[:] = ((status, headers),) - def raise_not_imp(*args): - raise NotImplementedError - return raise_not_imp + todo[:] = ((status, headers),) + def raise_not_imp(*args): + raise NotImplementedError + return raise_not_imp ret = self.app(environ, modified_start_response) assert hasattr(ret, "__iter__") - if todo and todo[0] is None: # nothing to do - #print "content-length: nothing" - return ret - if isinstance(ret, list): - #print "content-length: simple" + gotdata = True + assert bool(todo) status, headers = todo[0] - length = sum(map(len, ret)) - headers.append(("Content-length", str(length))) + if not [k for k, _ in headers if k.lower() == "content-length"]: + length = sum(map(len, ret)) + headers.append(("Content-Length", str(length))) start_response(status, headers) return ret ret = iter(ret) + first = "" stopped = False - data = CloseableList(getattr(ret, "close", None)) - length = 0 - try: - data.append(next(ret)) # fills todo - length += len(data[-1]) - except StopIteration: - stopped = True - + while not (first or stopped): + try: + first = next(ret) + except StopIteration: + stopped = True + gotdata = True + assert bool(todo) status, headers = todo[0] + data = CloseableList(getattr(ret, "close", None)) + if first: + data.append(first) + length = len(first) while (not stopped) and length < self.maxstore: try: @@ -191,12 +195,10 @@ class ContentLengthMiddleware: stopped = True if stopped: - #print "content-length: gathered" headers.append(("Content-length", str(length))) start_response(status, headers) return data - #print "content-length: passthrough" start_response(status, headers) return CloseableIterator(getattr(ret, "close", None), data, ret) @@ -276,10 +278,14 @@ class CachingMiddleware: del self.cache[path] cache_object = [now, "", [], []] def modified_start_respesponse(status, headers, exc_info=None): + try: + if cache_object[3]: + 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: - return self.app(status, headers, exc_info) cache_object[1] = status cache_object[2] = headers write = start_response(status, list(headers)) |