chore: enforce protocol type checking for app instance

Author: heitorlessa

aws_lambda_powertools/event_handler/api_gateway.py

@@ -9,26 +9,13 @@
 from enum import Enum
 from functools import partial
 from http import HTTPStatus
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    Match,
-    Optional,
-    Pattern,
-    Set,
-    Tuple,
-    Type,
-    Union,
-)
+from typing import Any, Callable, Dict, List, Match, Optional, Pattern, Set, Tuple, Type, Union
 
 from aws_lambda_powertools.event_handler import content_types
 from aws_lambda_powertools.event_handler.exceptions import NotFoundError, ServiceError
 from aws_lambda_powertools.shared.cookies import Cookie
 from aws_lambda_powertools.shared.functions import powertools_dev_is_set
 from aws_lambda_powertools.shared.json_encoder import Encoder
-from aws_lambda_powertools.shared.types import Protocol
 from aws_lambda_powertools.utilities.data_classes import (
     ALBEvent,
     APIGatewayProxyEvent,
@@ -1331,11 +1318,4 @@ def __init__(
         super().__init__(ProxyEventType.ALBEvent, cors, debug, serializer, strip_prefixes)
 
 
-class NextMiddlewareCallback(Protocol):
-    def __call__(self, app: ApiGatewayResolver) -> Response:
-        """Protocol for callback regardless of next_middleware(app), next(app)"""
-        ...
-
-    def __name__(self) -> str:  # noqa A003
-        """Protocol for name of the Middleware"""
-        ...
+# Specialized types defined here due to circular dependency in `.types`

aws_lambda_powertools/event_handler/middlewares/__init__.py

@@ -1,4 +1,4 @@
-from aws_lambda_powertools.event_handler.middlewares.base import BaseMiddlewareHandler
+from aws_lambda_powertools.event_handler.middlewares.base import BaseMiddlewareHandler, NextMiddleware
 from aws_lambda_powertools.event_handler.middlewares.schema_validation import SchemaValidationMiddleware
 
-__all__ = ["BaseMiddlewareHandler", "SchemaValidationMiddleware"]
+__all__ = ["BaseMiddlewareHandler", "SchemaValidationMiddleware", "NextMiddleware"]

aws_lambda_powertools/event_handler/middlewares/base.py

@@ -1,12 +1,23 @@
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Generic, TypeVar
+from typing import Generic
 
-from aws_lambda_powertools.event_handler.api_gateway import ApiGatewayResolver, Response
+from typing_extensions import Protocol
 
-EventHandlerResolver = TypeVar("EventHandlerResolver", bound=ApiGatewayResolver)
+from aws_lambda_powertools.event_handler.api_gateway import Response
+from aws_lambda_powertools.event_handler.types import EventHandlerInstance
 
 
-class BaseMiddlewareHandler(Generic[EventHandlerResolver], ABC):
+class NextMiddleware(Protocol):
+    def __call__(self, app: EventHandlerInstance) -> Response:
+        """Protocol for callback regardless of next_middleware(app), get_response(app) etc"""
+        ...
+
+    def __name__(self) -> str:  # noqa A003
+        """Protocol for name of the Middleware"""
+        ...
+
+
+class BaseMiddlewareHandler(Generic[EventHandlerInstance], ABC):
     """
     Base class for Middleware Handlers
 
@@ -69,15 +80,15 @@ class BaseMiddlewareHandler(Generic[EventHandlerResolver], ABC):
     """
 
     @abstractmethod
-    def handler(self, app: EventHandlerResolver, next_middleware: Callable[..., Any]) -> Response:
+    def handler(self, app: EventHandlerInstance, next_middleware: NextMiddleware) -> Response:
         """
         The Middleware Handler
 
         Parameters
         ----------
-        app: EventHandlerResolver
+        app: EventHandlerInstance
             An instance of an Event Handler that implements ApiGatewayResolver
-        next_middleware: Callable[..., Any]
+        next_middleware: NextMiddleware
             The next middleware handler in the chain
 
         Returns
@@ -92,7 +103,7 @@ def handler(self, app: EventHandlerResolver, next_middleware: Callable[..., Any]
     def __name__(self) -> str:  # noqa A003
         return str(self.__class__.__name__)
 
-    def __call__(self, app: EventHandlerResolver, next_middleware: Callable[..., Any]) -> Response:
+    def __call__(self, app: EventHandlerInstance, next_middleware: NextMiddleware) -> Response:
         """
         The Middleware handler function.
 

aws_lambda_powertools/event_handler/types.py

@@ -1,3 +1,5 @@
-from aws_lambda_powertools.event_handler.api_gateway import NextMiddlewareCallback
+from typing import TypeVar
 
-__all__ = ["NextMiddlewareCallback"]
+from aws_lambda_powertools.event_handler import ApiGatewayResolver
+
+EventHandlerInstance = TypeVar("EventHandlerInstance", bound=ApiGatewayResolver)

examples/event_handler_rest/src/custom_middlewares.py

@@ -1,12 +1,12 @@
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.event_handler import ApiGatewayResolver, Response
-from aws_lambda_powertools.event_handler.api_gateway import NextMiddlewareCallback
 from aws_lambda_powertools.event_handler.exceptions import BadRequestError, InternalServerError, ServiceError
+from aws_lambda_powertools.event_handler.middlewares import NextMiddleware
 
 logger = Logger()
 
 
-def validate_correlation_id(app: ApiGatewayResolver, next_middleware: NextMiddlewareCallback) -> Response:
+def validate_correlation_id(app: ApiGatewayResolver, next_middleware: NextMiddleware) -> Response:
     # If missing mandatory header raise an error
     if not app.current_event.headers.get("x-correlation-id", None):
         raise BadRequestError("No [x-correlation-id] header provided.  All requests must include this header.")
@@ -15,7 +15,7 @@ def validate_correlation_id(app: ApiGatewayResolver, next_middleware: NextMiddle
     return next_middleware(app)
 
 
-def sanitise_exceptions(app: ApiGatewayResolver, next_middleware: NextMiddlewareCallback) -> Response:
+def sanitise_exceptions(app: ApiGatewayResolver, next_middleware: NextMiddleware) -> Response:
     try:
         # Get the Result from the next middleware
         result = next_middleware(app)

examples/event_handler_rest/src/middleware_getting_started.py

@@ -2,13 +2,13 @@
 
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver, Response
-from aws_lambda_powertools.event_handler.types import NextMiddlewareCallback
+from aws_lambda_powertools.event_handler.middlewares import NextMiddleware
 
 app = APIGatewayRestResolver()
 logger = Logger()
 
 
-def inject_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMiddlewareCallback) -> Response:
+def inject_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMiddleware) -> Response:
     request_id = app.current_event.request_context.request_id  # (1)!
 
     # Use API Gateway REST API request ID if caller didn't include a correlation ID

examples/event_handler_rest/src/middleware_global_middlewares_module.py

@@ -1,11 +1,11 @@
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.event_handler import APIGatewayRestResolver, Response
-from aws_lambda_powertools.event_handler.types import NextMiddlewareCallback
+from aws_lambda_powertools.event_handler.middlewares import NextMiddleware
 
 logger = Logger()
 
 
-def log_request_response(app: APIGatewayRestResolver, next_middleware: NextMiddlewareCallback) -> Response:
+def log_request_response(app: APIGatewayRestResolver, next_middleware: NextMiddleware) -> Response:
     logger.info("Incoming request", path=app.current_event.path, request=app.current_event.raw_event)
 
     result = next_middleware(app)
@@ -14,7 +14,7 @@ def log_request_response(app: APIGatewayRestResolver, next_middleware: NextMiddl
     return result
 
 
-def inject_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMiddlewareCallback) -> Response:
+def inject_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMiddleware) -> Response:
     request_id = app.current_event.request_context.request_id
 
     # Use API Gateway REST API request ID if caller didn't include a correlation ID
@@ -32,7 +32,7 @@ def inject_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMidd
     return result
 
 
-def enforce_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMiddlewareCallback) -> Response:
+def enforce_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMiddleware) -> Response:
     # If missing mandatory header raise an error
     if not app.current_event.get_header_value("x-correlation-id", case_sensitive=False):
         return Response(status_code=400, body="Correlation ID header is now mandatory.")  # (1)!

tests/functional/event_handler/test_api_middlewares.py

@@ -6,13 +6,16 @@
     ApiGatewayResolver,
     APIGatewayRestResolver,
     BaseRouter,
-    NextMiddlewareCallback,
     ProxyEventType,
     Response,
     Router,
 )
 from aws_lambda_powertools.event_handler.exceptions import BadRequestError
-from aws_lambda_powertools.event_handler.middlewares import BaseMiddlewareHandler, SchemaValidationMiddleware
+from aws_lambda_powertools.event_handler.middlewares import (
+    BaseMiddlewareHandler,
+    NextMiddleware,
+    SchemaValidationMiddleware,
+)
 from tests.functional.utils import load_event
 
 API_REST_EVENT = load_event("apiGatewayProxyEvent.json")
@@ -29,15 +32,15 @@
 )
 def test_route_with_middleware(app: BaseRouter, event):
     # define custom middleware to inject new argument - "custom"
-    def middleware_1(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def middleware_1(app: BaseRouter, next_middleware: NextMiddleware):
         # add additional data to Router Context
         app.append_context(custom="custom")
         response = next_middleware(app)
 
         return response
 
     # define custom middleware to inject new argument - "another_one"
-    def middleware_2(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def middleware_2(app: BaseRouter, next_middleware: NextMiddleware):
         # add additional data to Router Context
         app.append_context(another_one=6)
         response = next_middleware(app)
@@ -84,15 +87,15 @@ def get_lambda() -> Response:
 )
 def test_with_router_middleware(app: BaseRouter, event, other_event):
     # define custom middleware to inject new argument - "custom"
-    def global_middleware(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def global_middleware(app: BaseRouter, next_middleware: NextMiddleware):
         # add custom data to context
         app.append_context(custom="custom")
         response = next_middleware(app)
 
         return response
 
     # define custom middleware to inject new argument - "another_one"
-    def middleware_2(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def middleware_2(app: BaseRouter, next_middleware: NextMiddleware):
         # add data to resolver context
         app.append_context(another_one=6)
         response = next_middleware(app)
@@ -143,7 +146,7 @@ def get_other_lambda() -> Response:
     ],
 )
 def test_dynamic_route_with_middleware(app: BaseRouter, event):
-    def middleware_one(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def middleware_one(app: BaseRouter, next_middleware: NextMiddleware):
         # inject data into the resolver context
         app.append_context(injected="injected_value")
         response = next_middleware(app)
@@ -182,12 +185,12 @@ def middleware_one(app: BaseRouter, next_middleware):
 
         return response
 
-    def early_return_middleware(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def early_return_middleware(app: BaseRouter, next_middleware: NextMiddleware):
         assert app.context.get("injected") == "injected_value"
 
         return Response(400, content_types.TEXT_HTML, "bad_response")
 
-    def not_executed_middleware(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def not_executed_middleware(app: BaseRouter, next_middleware: NextMiddleware):
         # This should never be executed - if it is an excpetion will be raised
         raise NotImplementedError()
 
@@ -310,19 +313,19 @@ def post_lambda() -> Response:
     ],
 )
 def test_middleware_short_circuit_via_httperrors(app: BaseRouter, event):
-    def middleware_one(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def middleware_one(app: BaseRouter, next_middleware: NextMiddleware):
         # inject a variable into the kwargs of the middleware chain
         app.append_context(injected="injected_value")
         response = next_middleware(app)
 
         return response
 
-    def early_return_middleware(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def early_return_middleware(app: BaseRouter, next_middleware: NextMiddleware):
         # ensure "injected" context variable is passed in by middleware_one
         assert app.context.get("injected") == "injected_value"
         raise BadRequestError("bad_response")
 
-    def not_executed_middleware(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def not_executed_middleware(app: BaseRouter, next_middleware: NextMiddleware):
         # This should never be executed - if it is an excpetion will be raised
         raise NotImplementedError()
 
@@ -373,7 +376,7 @@ def router_middleware(app: BaseRouter, next_middleware):
 
     to_inject: str = "injected_value"
 
-    def middleware_one(app: BaseRouter, next_middleware: NextMiddlewareCallback):
+    def middleware_one(app: BaseRouter, next_middleware: NextMiddleware):
         # inject a variable into the kwargs of the middleware chain
         app.append_context(injected=to_inject)
         response = next_middleware(app)
@@ -406,7 +409,7 @@ def __init__(self, header: str):
             super().__init__()
             self.header = header
 
-        def handler(self, app: ApiGatewayResolver, get_response: NextMiddlewareCallback, **kwargs) -> Response:
+        def handler(self, app: ApiGatewayResolver, get_response: NextMiddleware, **kwargs) -> Response:
             request_id = app.current_event.request_context.request_id  # type: ignore[attr-defined] # using REST event in a base Resolver # noqa: E501
             correlation_id = app.current_event.get_header_value(
                 name=self.header,