Simplify to just Callable[..., Any]

Author: huonw

aws_lambda_powertools/middleware_factory/factory.py

@@ -1,64 +1,19 @@
 import functools
 import inspect
 import logging
-import sys
 import os
-from typing import Any, Callable, Dict, Optional, Union, cast, overload
-
-if sys.version_info >= (3, 8):
-    from typing import Protocol
-else:
-    from typing_extensions import Protocol
+from typing import Any, Callable, Optional
 
 from ..shared import constants
 from ..shared.functions import resolve_truthy_env_var_choice
 from ..tracing import Tracer
-from ..utilities.typing import LambdaContext
 from .exceptions import MiddlewareInvalidArgumentError
 
 logger = logging.getLogger(__name__)
 
-# context: Any to avoid forcing users to type it as context: LambdaContext
-_Handler = Callable[[Any, LambdaContext], Any]
-_RawHandlerDecorator = Callable[[_Handler], _Handler]
-
-
-class _FactoryDecorator(Protocol):
-    # it'd be better for this to be using ParamSpec (available from 3.10)
-    def __call__(
-        self, handler: _Handler, event: Dict[str, Any], context: LambdaContext, **kwargs: Any
-    ) -> _RawHandlerDecorator:
-        ...
-
-
-class _HandlerDecorator(Protocol):
-    @overload
-    def __call__(self, decorator: _Handler) -> _Handler:
-        ...
-
-    @overload
-    def __call__(self, decorator: None = None, **kwargs: Any) -> _RawHandlerDecorator:
-        ...
-
-    def __call__(self, decorator: Optional[_Handler] = None, **kwargs: Any) -> Union[_Handler, _RawHandlerDecorator]:
-        ...
-
-
-@overload
-def lambda_handler_decorator(decorator: _FactoryDecorator) -> _HandlerDecorator:
-    ...
-
-
-@overload
-def lambda_handler_decorator(
-    decorator: None = None, trace_execution: Optional[bool] = None
-) -> Callable[[_FactoryDecorator], _HandlerDecorator]:
-    ...
-
 
-def lambda_handler_decorator(
-    decorator: Optional[_FactoryDecorator] = None, trace_execution: Optional[bool] = None
-) -> Union[_HandlerDecorator, Callable[[_FactoryDecorator], _HandlerDecorator]]:
+# giving this an accurate return type is hard
+def lambda_handler_decorator(decorator: Optional[Callable] = None, trace_execution: Optional[bool] = None) -> Callable:
     """Decorator factory for decorating Lambda handlers.
 
     You can use lambda_handler_decorator to create your own middlewares,
@@ -149,25 +104,19 @@ def lambda_handler(event, context):
     """
 
     if decorator is None:
-        return cast(
-            Callable[[_FactoryDecorator], _HandlerDecorator],
-            functools.partial(lambda_handler_decorator, trace_execution=trace_execution),
-        )
+        return functools.partial(lambda_handler_decorator, trace_execution=trace_execution)
 
     trace_execution = resolve_truthy_env_var_choice(
         env=os.getenv(constants.MIDDLEWARE_FACTORY_TRACE_ENV, "false"), choice=trace_execution
     )
 
     @functools.wraps(decorator)
-    def final_decorator(
-        func: Optional[_RawHandlerDecorator] = None, **kwargs: Any
-    ) -> Union[_Handler, _RawHandlerDecorator]:
+    def final_decorator(func: Optional[Callable] = None, **kwargs: Any):
         # If called with kwargs return new func with kwargs
         if func is None:
             return functools.partial(final_decorator, **kwargs)
 
         if not inspect.isfunction(func):
-            assert decorator is not None
             # @custom_middleware(True) vs @custom_middleware(log_event=True)
             raise MiddlewareInvalidArgumentError(
                 f"Only keyword arguments is supported for middlewares: {decorator.__qualname__} received {func}"  # type: ignore # noqa: E501
@@ -190,4 +139,4 @@ def wrapper(event, context):
 
         return wrapper
 
-    return cast(_HandlerDecorator, final_decorator)
+    return final_decorator