Use generics for logger.inject_lambda_context.

This enables static-type checking correctly with mypy when using
pydantic models for the lambda event argument rather than a mapping
type. Currently using this decorator results in an error and
requires a typing ignore directive on the line using the decorator.

Author: ascopes

aws_lambda_powertools/logging/logger.py

@@ -4,12 +4,13 @@
 import os
 import random
 import sys
-from typing import IO, Any, Callable, Dict, Iterable, Optional, TypeVar, Union
+from typing import IO, Any, Callable, Iterable, Optional, TypeVar, Union, overload
 
 import jmespath
 
 from ..shared import constants
 from ..shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
+from ..shared.types import AnyCallableT
 from .exceptions import InvalidLoggerSamplingRateError
 from .filters import SuppressFilter
 from .formatter import (
@@ -270,13 +271,33 @@ def _configure_sampling(self):
                 f"Please review POWERTOOLS_LOGGER_SAMPLE_RATE environment variable."
             )
 
+    @overload
     def inject_lambda_context(
         self,
-        lambda_handler: Optional[Callable[[Dict, Any], Any]] = None,
+        lambda_handler: AnyCallableT,
         log_event: Optional[bool] = None,
         correlation_id_path: Optional[str] = None,
         clear_state: Optional[bool] = False,
-    ):
+    ) -> AnyCallableT:
+        ...
+
+    @overload
+    def inject_lambda_context(
+        self,
+        lambda_handler: None = None,
+        log_event: Optional[bool] = None,
+        correlation_id_path: Optional[str] = None,
+        clear_state: Optional[bool] = False,
+    ) -> Callable[[AnyCallableT], AnyCallableT]:
+        ...
+
+    def inject_lambda_context(
+        self,
+        lambda_handler: Optional[AnyCallableT] = None,
+        log_event: Optional[bool] = None,
+        correlation_id_path: Optional[str] = None,
+        clear_state: Optional[bool] = False,
+    ) -> Any:
         """Decorator to capture Lambda contextual info and inject into logger
 
         Parameters