fix(typing): ensure return type is a T when default_value is set

Author: tonnico

aws_lambda_powertools/utilities/data_classes/api_gateway_authorizer_event.py

@@ -283,6 +283,17 @@ def path_parameters(self) -> Optional[Dict[str, str]]:
     def stage_variables(self) -> Optional[Dict[str, str]]:
         return self.get("stageVariables")
 
+    @overload
+    def get_header_value(self, name: str, default_value: str, case_sensitive: bool = False) -> str: ...
+
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: Optional[str] = None,
+        case_sensitive: Optional[bool] = False,
+    ) -> Optional[str]: ...
+
     def get_header_value(
         self,
         name: str,

aws_lambda_powertools/utilities/data_classes/appsync_resolver_event.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, overload
 
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 from aws_lambda_powertools.utilities.data_classes.shared_functions import (
@@ -214,6 +214,22 @@ def stash(self) -> Optional[dict]:
         a pipeline resolver."""
         return self.get("stash")
 
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: str,
+        case_sensitive: Optional[bool] = False,
+    ) -> str: ...
+
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: Optional[str] = None,
+        case_sensitive: Optional[bool] = False,
+    ) -> Optional[str]: ...
+
     def get_header_value(
         self,
         name: str,

aws_lambda_powertools/utilities/data_classes/common.py

@@ -2,7 +2,7 @@
 import json
 from collections.abc import Mapping
 from functools import cached_property
-from typing import Any, Callable, Dict, Iterator, List, Optional, overload
+from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar, overload
 
 from aws_lambda_powertools.shared.headers_serializer import BaseHeadersSerializer
 from aws_lambda_powertools.utilities.data_classes.shared_functions import (
@@ -11,6 +11,8 @@
     get_query_string_value,
 )
 
+T = TypeVar("T")
+
 
 class DictWrapper(Mapping):
     """Provides a single read only access to a wrapper dict"""
@@ -86,7 +88,13 @@ def _str_helper(self) -> Dict[str, Any]:
     def _properties(self) -> List[str]:
         return [p for p in dir(self.__class__) if isinstance(getattr(self.__class__, p), property)]
 
-    def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
+    @overload
+    def get(self, key: str, default: T) -> T: ...
+
+    @overload
+    def get(self, key: str, default: Optional[T] = None) -> Optional[T]: ...
+
+    def get(self, key: str, default: Optional[T] = None) -> Optional[T]:
         return self._data.get(key, default)
 
     @property
@@ -172,6 +180,12 @@ def http_method(self) -> str:
         """The HTTP method used. Valid values include: DELETE, GET, HEAD, OPTIONS, PATCH, POST, and PUT."""
         return self["httpMethod"]
 
+    @overload
+    def get_query_string_value(self, name: str, default_value: str) -> str: ...
+
+    @overload
+    def get_query_string_value(self, name: str, default_value: Optional[str] = None) -> Optional[str]: ...
+
     def get_query_string_value(self, name: str, default_value: Optional[str] = None) -> Optional[str]:
         """Get query string value by name
 

aws_lambda_powertools/utilities/data_classes/kafka_event.py

@@ -1,6 +1,6 @@
 import base64
 from functools import cached_property
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterator, List, Optional, overload
 
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 from aws_lambda_powertools.utilities.data_classes.shared_functions import (
@@ -69,10 +69,26 @@ def decoded_headers(self) -> Dict[str, bytes]:
         """Decodes the headers as a single dictionary."""
         return {k: bytes(v) for chunk in self.headers for k, v in chunk.items()}
 
+    @overload
     def get_header_value(
         self,
         name: str,
-        default_value: Optional[Any] = None,
+        default_value: str,
+        case_sensitive: bool = False,
+    ) -> str: ...
+
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: Optional[str] = None,
+        case_sensitive: bool = False,
+    ) -> Optional[str]: ...
+
+    def get_header_value(
+        self,
+        name: str,
+        default_value: Optional[str] = None,
         case_sensitive: bool = True,
     ) -> Optional[str]:
         """Get a decoded header value by name."""

aws_lambda_powertools/utilities/data_classes/s3_object_event.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, overload
 
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 from aws_lambda_powertools.utilities.data_classes.shared_functions import (
@@ -73,6 +73,22 @@ def headers(self) -> Dict[str, str]:
         The case of the original headers is retained in this map."""
         return self["headers"]
 
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: str,
+        case_sensitive: Optional[bool] = False,
+    ) -> str: ...
+
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: Optional[str] = None,
+        case_sensitive: Optional[bool] = False,
+    ) -> Optional[str]: ...
+
     def get_header_value(
         self,
         name: str,

aws_lambda_powertools/utilities/data_classes/shared_functions.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import base64
-from typing import Any, Dict
+from typing import Any, Dict, List, Optional, overload
 
 
 def base64_decode(value: str) -> str:
@@ -21,12 +21,30 @@ def base64_decode(value: str) -> str:
     return base64.b64decode(value).decode("UTF-8")
 
 
+@overload
 def get_header_value(
     headers: dict[str, Any],
     name: str,
-    default_value: str | None,
-    case_sensitive: bool | None,
-) -> str | None:
+    default_value: str,
+    case_sensitive: bool,
+) -> str: ...
+
+
+@overload
+def get_header_value(
+    headers: dict[str, Any],
+    name: str,
+    default_value: Optional[str],
+    case_sensitive: bool,
+) -> Optional[str]: ...
+
+
+def get_header_value(
+    headers: dict[str, Any],
+    name: str,
+    default_value: Optional[str],
+    case_sensitive: bool,
+) -> Optional[str]:
     """
     Get the value of a header by its name.
 
@@ -62,11 +80,27 @@ def get_header_value(
     )
 
 
+@overload
+def get_query_string_value(
+    query_string_parameters: Dict[str, str] | None,
+    name: str,
+    default_value: str,
+) -> str: ...
+
+
+@overload
+def get_query_string_value(
+    query_string_parameters: Dict[str, str] | None,
+    name: str,
+    default_value: Optional[str] = None,
+) -> Optional[str]: ...
+
+
 def get_query_string_value(
     query_string_parameters: Dict[str, str] | None,
     name: str,
-    default_value: str | None = None,
-) -> str | None:
+    default_value: Optional[str] = None,
+) -> Optional[str]:
     """
     Retrieves the value of a query string parameter specified by the given name.
 
@@ -87,18 +121,18 @@ def get_query_string_value(
 
 
 def get_multi_value_query_string_values(
-    multi_value_query_string_parameters: Dict[str, list[str]] | None,
+    multi_value_query_string_parameters: Dict[str, List[str]] | None,
     name: str,
-    default_values: list[str] | None = None,
-) -> list[str]:
+    default_values: Optional[List[str]] = None,
+) -> List[str]:
     """
     Retrieves the values of a multi-value string parameters specified by the given name.
 
     Parameters
     ----------
     name: str
         The name of the query string parameter to retrieve.
-    default_value: list[str], optional
+    default_value: List[str], optional
         The default value to return if the parameter is not found. Defaults to None.
 
     Returns

aws_lambda_powertools/utilities/data_classes/vpc_lattice.py

@@ -47,6 +47,12 @@ def http_method(self) -> str:
         """The HTTP method used. Valid values include: DELETE, GET, HEAD, OPTIONS, PATCH, POST, and PUT."""
         return self["method"]
 
+    @overload
+    def get_query_string_value(self, name: str, default_value: str) -> str: ...
+
+    @overload
+    def get_query_string_value(self, name: str, default_value: Optional[str] = None) -> Optional[str]: ...
+
     def get_query_string_value(self, name: str, default_value: Optional[str] = None) -> Optional[str]:
         """Get query string value by name