refactor(tracer): make capture_lambda_handler type more generic (#4796)

* chore(typing): update tracing capture_lambda_handler type

Signed-off-by: Amin Alaee <[email protected]>

* Making event argument generic and optional arguments Any

---------

Signed-off-by: Amin Alaee <[email protected]>
Co-authored-by: Leandro Damascena <[email protected]>
Co-authored-by: Simon Thulbourn <[email protected]>

Author: 3 people

aws_lambda_powertools/shared/types.py

@@ -19,6 +19,8 @@
     from typing_extensions import TypeAlias, get_args, get_origin
 
 AnyCallableT = TypeVar("AnyCallableT", bound=Callable[..., Any])  # noqa: VNE001
+
+
 # JSON primitives only, mypy doesn't support recursive tho
 JSONType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
 

aws_lambda_powertools/tracing/tracer.py

@@ -5,7 +5,7 @@
 import logging
 import numbers
 import os
-from typing import Any, Callable, Dict, List, Optional, Sequence, Union, cast, overload
+from typing import Any, Callable, Dict, List, Optional, Sequence, TypeVar, Union, cast, overload
 
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.shared.functions import (
@@ -22,6 +22,8 @@
 
 aws_xray_sdk = LazyLoader(constants.XRAY_SDK_MODULE, globals(), constants.XRAY_SDK_MODULE)
 
+T = TypeVar("T")
+
 
 class Tracer:
     """Tracer using AWS-XRay to provide decorators with known defaults for Lambda functions
@@ -250,7 +252,7 @@ def patch(self, modules: Optional[Sequence[str]] = None):
 
     def capture_lambda_handler(
         self,
-        lambda_handler: Union[Callable[[Dict, Any], Any], Optional[Callable[[Dict, Any, Optional[Dict]], Any]]] = None,
+        lambda_handler: Optional[Union[Callable[[T, Any], Any], Callable[[T, Any, Any], Any]]] = None,
         capture_response: Optional[bool] = None,
         capture_error: Optional[bool] = None,
     ):