feat: allow custom handlers to define their own cycle class via cycle_cls property

Author: khamaileon

docs/adapter.md

@@ -82,3 +82,101 @@ def hello(request: Request):
 
 handler = Mangum(app)
 ```
+
+## Custom Handlers
+
+Mangum supports custom handlers to process Lambda events that don't match the built-in handlers (API Gateway, ALB, Lambda@Edge). This is useful for handling custom event formats or integrating with other AWS services.
+
+### Basic Custom Handler
+
+A custom handler must implement the following interface:
+
+```python
+from mangum import Mangum
+from mangum.protocols import HTTPCycle
+from mangum.types import Cycle, LambdaConfig, LambdaContext, LambdaEvent, Response, Scope
+
+
+class MyCustomHandler:
+    @classmethod
+    def infer(cls, event: LambdaEvent, context: LambdaContext, config: LambdaConfig) -> bool:
+        """Return True if this handler should process the event."""
+        return "my-custom-key" in event
+
+    def __init__(self, event: LambdaEvent, context: LambdaContext, config: LambdaConfig) -> None:
+        self.event = event
+        self.context = context
+        self.config = config
+
+    @property
+    def cycle_cls(self) -> type[Cycle]:
+        """Return the cycle class to use for request/response processing."""
+        return HTTPCycle
+
+    @property
+    def body(self) -> bytes:
+        """Return the request body."""
+        return self.event.get("body", b"")
+
+    @property
+    def scope(self) -> Scope:
+        """Return the ASGI scope dictionary."""
+        return {
+            "type": "http",
+            "http_version": "1.1",
+            "method": self.event.get("method", "GET"),
+            "headers": [],
+            "path": self.event.get("path", "/"),
+            "raw_path": None,
+            "root_path": "",
+            "scheme": "https",
+            "query_string": b"",
+            "server": ("localhost", 443),
+            "client": ("127.0.0.1", 0),
+            "asgi": {"version": "3.0", "spec_version": "2.0"},
+            "aws.event": self.event,
+            "aws.context": self.context,
+        }
+
+    def __call__(self, response: Response) -> dict:
+        """Transform the ASGI response to a Lambda response."""
+        return {
+            "statusCode": response["status"],
+            "headers": {k.decode(): v.decode() for k, v in response["headers"]},
+            "body": response["body"].decode(),
+        }
+
+
+handler = Mangum(app, custom_handlers=[MyCustomHandler])
+```
+
+Custom handlers are checked **before** the built-in handlers, so they take priority.
+
+### Custom Protocol Cycle
+
+The `cycle_cls` property allows you to specify a custom request/response cycle class. This is useful for implementing custom protocols or adding middleware-like behavior at the cycle level.
+
+```python
+from mangum.protocols import HTTPCycle
+from mangum.types import ASGI, Cycle, Response, Scope
+
+
+class MyCustomCycle:
+    """A custom cycle that adds behavior to the standard HTTP cycle."""
+
+    def __init__(self, scope: Scope, body: bytes) -> None:
+        self.scope = scope
+        self._http_cycle = HTTPCycle(scope, body)
+
+    async def __call__(self, app: ASGI) -> Response:
+        # Add custom logic before/after the request
+        return await self._http_cycle(app)
+
+
+class MyCustomHandler:
+    # ... other methods ...
+
+    @property
+    def cycle_cls(self) -> type[Cycle]:
+        return MyCustomCycle
+```

mangum/adapter.py

@@ -7,7 +7,7 @@
 from mangum._compat import asyncio_run
 from mangum.exceptions import ConfigurationError
 from mangum.handlers import ALB, APIGateway, HTTPGateway, LambdaAtEdge
-from mangum.protocols import HTTPCycle, LifespanCycle
+from mangum.protocols import LifespanCycle
 from mangum.types import ASGI, LambdaConfig, LambdaContext, LambdaEvent, LambdaHandler, LifespanMode
 
 logger = logging.getLogger("mangum")
@@ -67,12 +67,12 @@ async def handle_request() -> dict[str, Any]:
                 lifespan_cycle = LifespanCycle(self.app, self.lifespan)
                 async with lifespan_cycle:
                     scope.update({"state": lifespan_cycle.lifespan_state.copy()})
-                    http_cycle = HTTPCycle(scope, handler.body)
-                    http_response = await http_cycle(self.app)
-                    return handler(http_response)
+                    cycle = handler.cycle_cls(scope, handler.body)
+                    response = await cycle(self.app)
+                    return handler(response)
             else:
-                http_cycle = HTTPCycle(scope, handler.body)
-                http_response = await http_cycle(self.app)
-                return handler(http_response)
+                cycle = handler.cycle_cls(scope, handler.body)
+                response = await cycle(self.app)
+                return handler(response)
 
         return asyncio_run(handle_request())

mangum/handlers/alb.py

@@ -11,7 +11,9 @@
     handle_exclude_headers,
     maybe_encode_body,
 )
+from mangum.protocols import HTTPCycle
 from mangum.types import (
+    Cycle,
     LambdaConfig,
     LambdaContext,
     LambdaEvent,
@@ -94,6 +96,10 @@ def __init__(self, event: LambdaEvent, context: LambdaContext, config: LambdaCon
         self.context = context
         self.config = config
 
+    @property
+    def cycle_cls(self) -> type[Cycle]:
+        return HTTPCycle
+
     @property
     def body(self) -> bytes:
         return maybe_encode_body(

mangum/handlers/api_gateway.py

@@ -11,7 +11,9 @@
     maybe_encode_body,
     strip_api_gateway_path,
 )
+from mangum.protocols import HTTPCycle
 from mangum.types import (
+    Cycle,
     Headers,
     LambdaConfig,
     LambdaContext,
@@ -74,6 +76,10 @@ def __init__(self, event: LambdaEvent, context: LambdaContext, config: LambdaCon
         self.context = context
         self.config = config
 
+    @property
+    def cycle_cls(self) -> type[Cycle]:
+        return HTTPCycle
+
     @property
     def body(self) -> bytes:
         return maybe_encode_body(
@@ -132,6 +138,10 @@ def __init__(self, event: LambdaEvent, context: LambdaContext, config: LambdaCon
         self.context = context
         self.config = config
 
+    @property
+    def cycle_cls(self) -> type[Cycle]:
+        return HTTPCycle
+
     @property
     def body(self) -> bytes:
         return maybe_encode_body(

mangum/handlers/lambda_at_edge.py

@@ -8,7 +8,15 @@
     handle_multi_value_headers,
     maybe_encode_body,
 )
-from mangum.types import LambdaConfig, LambdaContext, LambdaEvent, Response, Scope
+from mangum.protocols import HTTPCycle
+from mangum.types import (
+    Cycle,
+    LambdaConfig,
+    LambdaContext,
+    LambdaEvent,
+    Response,
+    Scope,
+)
 
 
 class LambdaAtEdge:
@@ -25,6 +33,10 @@ def __init__(self, event: LambdaEvent, context: LambdaContext, config: LambdaCon
         self.context = context
         self.config = config
 
+    @property
+    def cycle_cls(self) -> type[Cycle]:
+        return HTTPCycle
+
     @property
     def body(self) -> bytes:
         cf_request_body = self.event["Records"][0]["cf"]["request"].get("body", {})

mangum/types.py

@@ -110,12 +110,27 @@ class LambdaConfig(TypedDict):
     exclude_headers: list[str]
 
 
+class Cycle(Protocol):
+    def __init__(self, scope: Scope, body: bytes) -> None: ...  # pragma: no cover
+
+    async def __call__(self, app: ASGI) -> Response: ...  # pragma: no cover
+
+    async def run(self, app: ASGI) -> None: ...  # pragma: no cover
+
+    async def receive(self) -> Message: ...  # pragma: no cover
+
+    async def send(self, message: Message) -> None: ...  # pragma: no cover
+
+
 class LambdaHandler(Protocol):
     def __init__(self, *args: Any) -> None: ...  # pragma: no cover
 
     @classmethod
     def infer(cls, event: LambdaEvent, context: LambdaContext, config: LambdaConfig) -> bool: ...  # pragma: no cover
 
+    @property
+    def cycle_cls(self) -> type[Cycle]: ...  # pragma: no cover
+
     @property
     def body(self) -> bytes: ...  # pragma: no cover
 

tests/handlers/test_custom.py

@@ -81,3 +81,60 @@ def test_custom_handler_call():
 
     result = handler(status=200, headers=[], body=b"Hello, World!")
     assert result == {"statusCode": 200, "headers": {}, "body": "Hello, World!"}
+
+
+def test_custom_handler_with_custom_cycle_cls():
+    """Test that a handler can define a custom cycle_cls property."""
+    from mangum import Mangum
+    from mangum.protocols import HTTPCycle
+    from mangum.types import Cycle, Response
+
+    class HandlerWithCustomCycle:
+        @classmethod
+        def infer(cls, event, context, config):
+            return "custom-cycle" in event
+
+        def __init__(self, event, context, config):
+            self.event = event
+            self.context = context
+            self.config = config
+
+        @property
+        def cycle_cls(self) -> type[Cycle]:
+            return HTTPCycle
+
+        @property
+        def body(self) -> bytes:
+            return b""
+
+        @property
+        def scope(self) -> Scope:
+            return {
+                "type": "http",
+                "http_version": "1.1",
+                "method": "GET",
+                "headers": [],
+                "path": "/",
+                "raw_path": None,
+                "root_path": "",
+                "scheme": "https",
+                "query_string": b"",
+                "server": ("localhost", 8080),
+                "client": ("127.0.0.1", 0),
+                "asgi": {"version": "3.0", "spec_version": "2.0"},
+                "aws.event": self.event,
+                "aws.context": self.context,
+            }
+
+        def __call__(self, response: Response) -> dict:
+            return {"statusCode": response["status"], "body": response["body"].decode()}
+
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"OK"})
+
+    handler = Mangum(app, lifespan="off", custom_handlers=[HandlerWithCustomCycle])
+    response = handler({"custom-cycle": True}, {})
+
+    assert response["statusCode"] == 200
+    assert response["body"] == "OK"