Refactor 204 handling.

Place the HTTP 204 handling in a `finally` at the bottom of the WSGI __call__
to ensure we catch *all* changes to `state.response`.

Change-Id: Id150205f295a3c0ca9c73ad897a41d4790476811

1 files changed, 20 insertions, 14 deletions

diff --git a/pecan/core.py b/pecan/core.py
index cb4cf2c..29de33a 100644
--- a/pecan/core.py
+++ b/pecan/core.py
@@ -566,33 +566,37 @@ class PecanBase(object):
             resp.text = result
         elif result:
             resp.body = result
-        elif response.status_int == 200:
+
+        if pecan_state['content_type']:
+            # set the content type
+            resp.content_type = pecan_state['content_type']
+
+    def _handle_empty_response_body(self, state):
+        # Enforce HTTP 204 for responses which contain no body
+        if state.response.status_int == 200:
             # If the response is a generator...
-            if isinstance(response.app_iter, types.GeneratorType):
+            if isinstance(state.response.app_iter, types.GeneratorType):
                 # Split the generator into two so we can peek at one of them
                 # and determine if there is any response body content
-                a, b = tee(response.app_iter)
+                a, b = tee(state.response.app_iter)
                 try:
                     next(a)
                 except StopIteration:
                     # If we hit StopIteration, the body is empty
-                    resp.status = 204
+                    state.response.status = 204
                 finally:
-                    resp.app_iter = b
+                    state.response.app_iter = b
             else:
                 text = None
-                if response.charset:
+                if state.response.charset:
                     # `response.text` cannot be accessed without a charset
                     # (because we don't know which encoding to use)
-                    text = response.text
-                if not any((response.body, text)):
-                    resp.status = 204
+                    text = state.response.text
+                if not any((state.response.body, text)):
+                    state.response.status = 204
 
-        if resp.status_int in (204, 304):
-            resp.content_type = None
-        elif pecan_state['content_type']:
-            # set the content type
-            resp.content_type = pecan_state['content_type']
+        if state.response.status_int in (204, 304):
+            state.response.content_type = None
 
     def __call__(self, environ, start_response):
         '''
@@ -663,6 +667,8 @@ class PecanBase(object):
                 self.determine_hooks(state.controller), 'after', state
             )
 
+        self._handle_empty_response_body(state)
+
         # get the response
         return state.response(environ, start_response)