From 6eb8443dcbe08e5a2c2be732fe1f9e4793f1dd42 Mon Sep 17 00:00:00 2001
From: Ashley Scopes <73482956+ascopes@users.noreply.github.com>
Date: Tue, 11 Oct 2022 09:15:26 +0100
Subject: [PATCH] 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.
---
 aws_lambda_powertools/logging/logger.py | 27 ++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py
index a0e24f1dcf9..0b2932ee684 100644
--- a/aws_lambda_powertools/logging/logger.py
+++ b/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