From e47af7ae7d3a78d86981e7fc3a583e23893973bc Mon Sep 17 00:00:00 2001
From: Bert Blommers <info@bertblommers.nl>
Date: Thu, 24 Jul 2025 10:58:03 +0000
Subject: [PATCH] Introduce a  to the RouterHandler

---
 rolo/gateway/handlers.py       | 14 +++++++++++++-
 tests/gateway/test_handlers.py | 27 +++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/rolo/gateway/handlers.py b/rolo/gateway/handlers.py
index 43b8cdb..ce4324a 100644
--- a/rolo/gateway/handlers.py
+++ b/rolo/gateway/handlers.py
@@ -18,13 +18,25 @@ class RouterHandler:
     router: Router
     respond_not_found: bool
 
-    def __init__(self, router: Router, respond_not_found: bool = False) -> None:
+    def __init__(
+        self,
+        router: Router,
+        respond_not_found: bool = False,
+        response_converter: t.Optional[t.Callable[..., Response]] = None,
+    ) -> None:
+        """
+        Individual routes from the Router should return a `rolo.response.Response` object.
+        If routes (can) return a custom response, use the `response_converter`-argument to pass a function that convert that custom response to a `rolo.response.Response` object.
+        """
         self.router = router
         self.respond_not_found = respond_not_found
+        self.response_converter = response_converter
 
     def __call__(self, chain: HandlerChain, context: RequestContext, response: Response):
         try:
             router_response = self.router.dispatch(context.request)
+            if self.response_converter:
+                router_response = self.response_converter(router_response)
             response.update_from(router_response)
             chain.stop()
         except NotFound:
diff --git a/tests/gateway/test_handlers.py b/tests/gateway/test_handlers.py
index 68f80cd..121da99 100644
--- a/tests/gateway/test_handlers.py
+++ b/tests/gateway/test_handlers.py
@@ -70,6 +70,33 @@ def test_router_handler_with_respond_not_found(self, serve_gateway):
         assert response.status_code == 404
         assert response.text == "not found"
 
+    def test_router_handler_with_custom_object(self, serve_gateway):
+        class CustomObject:
+            def __init__(self, message: str):
+                self.message = message
+
+        def _custom_object_handler(request: Request, args) -> CustomObject:
+            return (200, CustomObject(message=request.path))
+
+        def _response_converter(response: tuple[int, CustomObject]) -> Response:
+            status, body = response
+            return Response(status=status, response=body.message)
+
+        router = Router()
+        router.add("/foo", _custom_object_handler)
+
+        server = serve_gateway(
+            Gateway(
+                request_handlers=[
+                    RouterHandler(router, response_converter=_response_converter),
+                ],
+            )
+        )
+
+        doc = requests.get(server.url + "/foo")
+        assert doc.status_code == 200
+        assert doc.content == b"/foo"
+
 
 @pytest.mark.parametrize("serve_gateway", ["wsgi", "asgi", "twisted"], indirect=True)
 class TestEmptyResponseHandler: