Added a middleware to clear out error bodies, making them more accessible to Apache; added an app_iter wrapper for chaining app_iters from multiple sources (needed for peeking at status)

4 files changed, 110 insertions, 3 deletions

diff --git a/paste/errordocument.py b/paste/errordocument.py
index 991ba8a..0a44a3e 100644
--- a/paste/errordocument.py
+++ b/paste/errordocument.py
@@ -59,6 +59,7 @@ base path of the application is.
 
 from urllib import urlencode
 from urlparse import urlparse
+from paste.wsgilib import chained_app_iters
 
 def forward(app, codes):
     """
@@ -327,3 +328,51 @@ def make_errordocument(app, global_conf, **kw):
         map[status] = redir_loc
     forwarder = forward(app, map)
     return forwarder
+
+
+def make_empty_error(app, global_conf, **kw):
+    """
+    Use like:
+
+      [filter-app:main]
+      use = egg:Paste#emptyerror
+      next = real-app
+
+    This will clear the body of any bad responses (e.g., 404, 500,
+    etc).  If running behind Apache, Apache will replace the empty
+    response with whatever its configured ``ErrorDocument`` (but
+    Apache doesn't overwrite responses that do have content, which is
+    why this middlware is necessary)
+    """
+    if kw:
+        raise ValueError(
+            'emptyerror does not take any configuration')
+    return empty_error(app)
+
+def empty_error(app):
+    def filtered_app(environ, start_response):
+        got_status = []
+        def replace_start_response(status, headers, exc_info=None):
+            got_status.append(status)
+            return start_response(status, headers, exc_info)
+        app_iter = app(environ, replace_start_response)
+        item1 = None
+        if not got_status:
+            item1 = ''
+            for item in app_iter:
+                item1 = item
+                break
+            if not got_status:
+                raise ValueError(
+                    "start_response not called from application")
+        status = int(got_status[0].split()[0])
+        if status >= 400:
+            if hasattr(app_iter, 'close'):
+                app_iter.close()
+            return ['']
+        else:
+            if item1 is not None:
+                return chained_app_iters([item1], app_iter)
+            else:
+                return app_iter
+    return filtered_app
diff --git a/paste/wsgilib.py b/paste/wsgilib.py
index 85c646a..59f63ab 100644
--- a/paste/wsgilib.py
+++ b/paste/wsgilib.py
@@ -27,7 +27,8 @@ __all__ = ['get_cookies', 'add_close', 'raw_interactive',
            'interactive', 'construct_url', 'error_body_response',
            'error_response', 'send_file', 'has_header', 'header_value',
            'path_info_split', 'path_info_pop', 'capture_output',
-           'catch_errors', 'dump_environ', 'intercept_output']
+           'catch_errors', 'dump_environ', 'intercept_output',
+           'chained_app_iters']
 
 
 class add_close:
@@ -101,6 +102,51 @@ class add_start_close:
                 "WSGI request.  finalization function %s not called"
                 % self.close_func)
 
+class chained_app_iters:
+
+    """
+    Chains several app_iters together, also delegating .close() to each
+    of them.
+    """
+
+    def __init__(self, *chained):
+        self.app_iters = chained
+        self.chained = [iter(item) for item in chained]
+        self._closed = False
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if len(self.chained) == 1:
+            return self.chained[0].next()
+        else:
+            try:
+                return self.chained[0].next()
+            except StopIteration:
+                self.chained.pop(0)
+                return self.next()
+
+    def close(self):
+        self._closed = True
+        got_exc = None
+        for app_iter in self.app_iters:
+            try:
+                if hasattr(app_iter, 'close'):
+                    app_iter.close()
+            except:
+                got_exc = sys.exc_info()
+        if got_exc:
+            raise got_exc[0], got_exc[1], got_exc[2]
+
+    def __del__(self):
+        if not self._closed:
+            # We can't raise an error or anything at this stage
+            print >> sys.stderr, (
+                "Error: app_iter.close() was not called when finishing "
+                "WSGI request.  finalization function %s not called"
+                % self.close_func)
+
 def catch_errors(application, environ, start_response, error_callback,
                  ok_callback=None):
     """
diff --git a/setup.py b/setup.py
index 997e620..0fcfbf8 100644
--- a/setup.py
+++ b/setup.py
@@ -174,6 +174,7 @@ For the latest changes see the `news file
       recorder = paste.debug.recorder.record:make_recorder
       pony = paste.pony:make_pony
       errordocument = paste.errordocument:make_errordocument
+      emptyerror = paste.errordocument.make_empty_error
 
       [paste.server_runner]
       http = paste.httpserver:server_runner
diff --git a/tests/test_errordocument.py b/tests/test_errordocument.py
index fab7cca..51135b6 100644
--- a/tests/test_errordocument.py
+++ b/tests/test_errordocument.py
@@ -14,8 +14,7 @@ I also need to find out how to test that another response was
 correctly requested by the middleware.
 """
 import os
-import py.test
-from paste.errordocument import forward, custom_forward
+from paste.errordocument import forward, custom_forward, empty_error
 from paste.fixture import *
 
 def simple_app(environ, start_response):
@@ -32,3 +31,15 @@ def test_ok():
     assert res.header('content-type') == 'text/plain'
     assert res.full_status == '200 OK'
     assert 'requested page returned' in res
+
+def test_empty():
+    app = TestApp(empty_error(simple_app))
+    res = app.get('/')
+    assert res.header('content-type') == 'text/plain'
+    assert res.full_status == '200 OK'
+    assert 'requested page returned' in res
+    app = TestApp(empty_error(not_found_app))
+    res = app.get('/', status=404)
+    assert 'requested page returned' not in res
+    assert res.body == ''
+    assert res.status == 404