From cedb5c9a1529f2ebbf26544276159f9f7289ebdb Mon Sep 17 00:00:00 2001
From: Troy Swanson <1420926+troyswanson@users.noreply.github.com>
Date: Wed, 20 Dec 2023 10:05:18 -0600
Subject: [PATCH] refactor(parameters): add overload signatures for
 get_parameter and get_parameters (#3534)

Co-authored-by: Ruben Fonseca <rubefons@amazon.com>
---
 .../utilities/parameters/ssm.py               | 108 +++++++++++++++++-
 .../getting_started_single_ssm_parameter.py   |   2 +-
 2 files changed, 107 insertions(+), 3 deletions(-)

diff --git a/aws_lambda_powertools/utilities/parameters/ssm.py b/aws_lambda_powertools/utilities/parameters/ssm.py
index ebe8e57ac2..7df0e66c02 100644
--- a/aws_lambda_powertools/utilities/parameters/ssm.py
+++ b/aws_lambda_powertools/utilities/parameters/ssm.py
@@ -545,9 +545,57 @@ def _raise_if_errors_key_is_present(parameters: Dict, reserved_parameter: str, r
             )
 
 
+@overload
+def get_parameter(
+    name: str,
+    transform: None = None,
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> str:
+    ...
+
+
+@overload
+def get_parameter(
+    name: str,
+    transform: Literal["json"],
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> dict:
+    ...
+
+
+@overload
 def get_parameter(
     name: str,
-    transform: Optional[str] = None,
+    transform: Literal["binary"],
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> Union[str, dict, bytes]:
+    ...
+
+
+@overload
+def get_parameter(
+    name: str,
+    transform: Literal["auto"],
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> bytes:
+    ...
+
+
+def get_parameter(
+    name: str,
+    transform: TransformOptions = None,
     decrypt: Optional[bool] = None,
     force_fetch: bool = False,
     max_age: Optional[int] = None,
@@ -625,9 +673,65 @@ def get_parameter(
     )
 
 
+@overload
 def get_parameters(
     path: str,
-    transform: Optional[str] = None,
+    transform: None = None,
+    recursive: bool = True,
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    raise_on_transform_error: bool = False,
+    **sdk_options,
+) -> Dict[str, str]:
+    ...
+
+
+@overload
+def get_parameters(
+    path: str,
+    transform: Literal["json"],
+    recursive: bool = True,
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    raise_on_transform_error: bool = False,
+    **sdk_options,
+) -> Dict[str, dict]:
+    ...
+
+
+@overload
+def get_parameters(
+    path: str,
+    transform: Literal["binary"],
+    recursive: bool = True,
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    raise_on_transform_error: bool = False,
+    **sdk_options,
+) -> Dict[str, bytes]:
+    ...
+
+
+@overload
+def get_parameters(
+    path: str,
+    transform: Literal["auto"],
+    recursive: bool = True,
+    decrypt: Optional[bool] = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    raise_on_transform_error: bool = False,
+    **sdk_options,
+) -> Union[Dict[str, bytes], Dict[str, dict], Dict[str, str]]:
+    ...
+
+
+def get_parameters(
+    path: str,
+    transform: TransformOptions = None,
     recursive: bool = True,
     decrypt: Optional[bool] = None,
     force_fetch: bool = False,
diff --git a/examples/parameters/src/getting_started_single_ssm_parameter.py b/examples/parameters/src/getting_started_single_ssm_parameter.py
index d31c7a180f..f35aa3ede8 100644
--- a/examples/parameters/src/getting_started_single_ssm_parameter.py
+++ b/examples/parameters/src/getting_started_single_ssm_parameter.py
@@ -7,7 +7,7 @@
 def lambda_handler(event: dict, context: LambdaContext) -> dict:
     try:
         # Retrieve a single parameter
-        endpoint_comments: str = parameters.get_parameter("/lambda-powertools/endpoint_comments")  # type: ignore[assignment] # noqa: E501
+        endpoint_comments = parameters.get_parameter("/lambda-powertools/endpoint_comments")
 
         # the value of this parameter is https://jsonplaceholder.typicode.com/comments/
         comments: requests.Response = requests.get(endpoint_comments)