Added another app_iter wrapper similar to add_close, except that also calls a function just before the first content is returned

1 files changed, 39 insertions, 0 deletions

diff --git a/paste/wsgilib.py b/paste/wsgilib.py
index 1637db3..85c646a 100644
--- a/paste/wsgilib.py
+++ b/paste/wsgilib.py
@@ -62,6 +62,45 @@ class add_close:
                 "WSGI request.  finalization function %s not called"
                 % self.close_func)
 
+class add_start_close:
+    """
+    An an iterable that iterates over app_iter, calls start_func
+    before the first item is returned, then calls close_func at the
+    end.
+    """
+
+    def __init__(self, app_iterable, start_func, close_func=None):
+        self.app_iterable = app_iterable
+        self.app_iter = iter(app_iterable)
+        self.first = True
+        self.start_func = start_func
+        self.close_func = close_func
+        self._closed = False
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if self.first:
+            self.start_func()
+            self.first = False
+        return self.app_iter.next()
+
+    def close(self):
+        self._closed = True
+        if hasattr(self.app_iterable, 'close'):
+            self.app_iterable.close()
+        if self.close_func is not None:
+            self.close_func()
+
+    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):
     """