Add @after_error_handler decorator (Fixes #97)

docs/intro.rst

@@ -293,6 +293,12 @@ The function can return a modified response object to replace the original. If
 the function does not return a value, then the original response object is
 used.
 
+The after request handlers are only invoked for successful requests. The
+:func:`after_error_request() <microdot.Microdot.after_error_request>`
+decorator can be used to register a function that is called after an error
+occurs. The function receives the request and the error response and is
+expected to return an updated response object.
+
 .. note::
    The :ref:`request.g <The "g" Object>` object is a special object that allows
    the before and after request handlers, as well sa the route function to

src/microdot.py

@@ -472,6 +472,9 @@ class Request():
                     return response
 
                 return 'Hello, World!'
+
+        Note that the function is not called if the request handler raises an
+        exception and an error response is returned instead.
         """
         self.after_request_handlers.append(f)
         return f
@@ -746,6 +749,7 @@ class Microdot():
         self.url_map = []
         self.before_request_handlers = []
         self.after_request_handlers = []
+        self.after_error_request_handlers = []
         self.error_handlers = {}
         self.shutdown_requested = False
         self.debug = False
@@ -907,6 +911,24 @@ class Microdot():
         self.after_request_handlers.append(f)
         return f
 
+    def after_error_request(self, f):
+        """Decorator to register a function to run after an error response is
+        generated. The decorated function must take two arguments, the request
+        and response objects. The return value of the function must be an
+        updated response object. The handler is invoked for error responses
+        generated by Microdot, as well as those returned by application-defined
+        error handlers.
+
+        Example::
+
+            @app.after_error_request
+            def func(request, response):
+                # ...
+                return response
+        """
+        self.after_error_request_handlers.append(f)
+        return f
+
     def errorhandler(self, status_code_or_exception_class):
         """Decorator to register a function as an error handler. Error handler
         functions for numeric HTTP status codes must accept a single argument,
@@ -947,6 +969,8 @@ class Microdot():
             self.before_request_handlers.append(handler)
         for handler in subapp.after_request_handlers:
             self.after_request_handlers.append(handler)
+        for handler in subapp.after_error_request_handlers:
+            self.after_error_request_handlers.append(handler)
         for status_code, handler in subapp.error_handlers.items():
             self.error_handlers[status_code] = handler
 
@@ -1094,6 +1118,7 @@ class Microdot():
                 status_code=res.status_code))
 
     def dispatch_request(self, req):
+        after_request_handled = False
         if req:
             if req.content_length > req.max_content_length:
                 if 413 in self.error_handlers:
@@ -1126,6 +1151,7 @@ class Microdot():
                             res = handler(req, res) or res
                         for handler in req.after_request_handlers:
                             res = handler(req, res) or res
+                        after_request_handled = True
                     elif f in self.error_handlers:
                         res = self.error_handlers[f](req)
                     else:
@@ -1166,6 +1192,9 @@ class Microdot():
             res = Response(*res)
         elif not isinstance(res, Response):
             res = Response(res)
+        if not after_request_handled:
+            for handler in self.after_error_request_handlers:
+                res = handler(req, res) or res
         return res
 
 

src/microdot_asyncio.py

@@ -347,6 +347,7 @@ class Microdot(BaseMicrodot):
                 status_code=res.status_code))
 
     async def dispatch_request(self, req):
+        after_request_handled = False
         if req:
             if req.content_length > req.max_content_length:
                 if 413 in self.error_handlers:
@@ -383,6 +384,7 @@ class Microdot(BaseMicrodot):
                         for handler in req.after_request_handlers:
                             res = await self._invoke_handler(
                                 handler, req, res) or res
+                        after_request_handled = True
                     elif f in self.error_handlers:
                         res = await self._invoke_handler(
                             self.error_handlers[f], req)
@@ -425,6 +427,10 @@ class Microdot(BaseMicrodot):
             res = Response(*res)
         elif not isinstance(res, Response):
             res = Response(res)
+        if not after_request_handled:
+            for handler in self.after_error_request_handlers:
+                res = await self._invoke_handler(
+                    handler, req, res) or res
         return res
 
     async def _invoke_handler(self, f_or_coro, *args, **kwargs):

tests/test_microdot.py

@@ -279,6 +279,39 @@ class TestMicrodot(unittest.TestCase):
         self.assertEqual(res.headers['Content-Length'], '3')
         self.assertEqual(res.text, 'baz')
 
+    def test_after_error_request(self):
+        app = Microdot()
+
+        @app.after_error_request
+        def after_error_request_one(req, res):
+            res.headers['X-One'] = '1'
+
+        @app.after_error_request
+        def after_error_request_two(req, res):
+            res.set_cookie('foo', 'bar')
+            return res
+
+        @app.route('/foo')
+        def foo(req):
+            return 'foo'
+
+        client = TestClient(app)
+
+        res = client.get('/foo')
+        self.assertEqual(res.status_code, 200)
+        self.assertEqual(res.headers['Content-Type'],
+                         'text/plain; charset=UTF-8')
+        self.assertNotIn('X-One', res.headers)
+        self.assertNotIn('Set-Cookie', res.headers)
+
+        res = client.get('/bar')
+        self.assertEqual(res.status_code, 404)
+        self.assertEqual(res.headers['Content-Type'],
+                         'text/plain; charset=UTF-8')
+        self.assertEqual(res.headers['Set-Cookie'], ['foo=bar'])
+        self.assertEqual(res.headers['X-One'], '1')
+        self.assertEqual(client.cookies['foo'], 'bar')
+
     def test_400(self):
         self._mock()
 
@@ -661,7 +694,11 @@ class TestMicrodot(unittest.TestCase):
 
         @subapp.after_request
         def after(req, res):
-            return res.body + b':after'
+            res.body += b':after'
+
+        @subapp.after_error_request
+        def after_error(req, res):
+            res.body += b':errorafter'
 
         @subapp.errorhandler(404)
         def not_found(req):
@@ -680,7 +717,7 @@ class TestMicrodot(unittest.TestCase):
         self.assertEqual(res.status_code, 404)
         self.assertEqual(res.headers['Content-Type'],
                          'text/plain; charset=UTF-8')
-        self.assertEqual(res.text, '404')
+        self.assertEqual(res.text, '404:errorafter')
 
         res = client.get('/sub/app')
         self.assertEqual(res.status_code, 200)

tests/test_microdot_asyncio.py

@@ -314,6 +314,39 @@ class TestMicrodotAsync(unittest.TestCase):
         self.assertEqual(res.headers['Content-Length'], '3')
         self.assertEqual(res.text, 'baz')
 
+    def test_after_error_request(self):
+        app = Microdot()
+
+        @app.after_error_request
+        def after_error_request_one(req, res):
+            res.headers['X-One'] = '1'
+
+        @app.after_error_request
+        def after_error_request_two(req, res):
+            res.set_cookie('foo', 'bar')
+            return res
+
+        @app.route('/foo')
+        def foo(req):
+            return 'foo'
+
+        client = TestClient(app)
+
+        res = self._run(client.get('/foo'))
+        self.assertEqual(res.status_code, 200)
+        self.assertEqual(res.headers['Content-Type'],
+                         'text/plain; charset=UTF-8')
+        self.assertNotIn('X-One', res.headers)
+        self.assertNotIn('Set-Cookie', res.headers)
+
+        res = self._run(client.get('/bar'))
+        self.assertEqual(res.status_code, 404)
+        self.assertEqual(res.headers['Content-Type'],
+                         'text/plain; charset=UTF-8')
+        self.assertEqual(res.headers['Set-Cookie'], ['foo=bar'])
+        self.assertEqual(res.headers['X-One'], '1')
+        self.assertEqual(client.cookies['foo'], 'bar')
+
     def test_400(self):
         self._mock()