From 1b4a57ee75f9b648daf43b60be23f38e41219149 Mon Sep 17 00:00:00 2001
From: Tony Sherman <100969281+TonySherman@users.noreply.github.com>
Date: Fri, 8 Mar 2024 10:26:05 -0500
Subject: [PATCH] refactor(parameters): improve typing for get_secret method
 (#3910)

fix: add overloads for get_secret function

Co-authored-by: Leandro Damascena <lcdama@amazon.pt>
---
 .../utilities/parameters/secrets.py           | 46 ++++++++++++++++++-
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/aws_lambda_powertools/utilities/parameters/secrets.py b/aws_lambda_powertools/utilities/parameters/secrets.py
index beb4bb80846..9b09367cadb 100644
--- a/aws_lambda_powertools/utilities/parameters/secrets.py
+++ b/aws_lambda_powertools/utilities/parameters/secrets.py
@@ -3,11 +3,13 @@
 """
 
 import os
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Union, overload
 
 import boto3
 from botocore.config import Config
 
+from aws_lambda_powertools.utilities.parameters.types import TransformOptions
+
 if TYPE_CHECKING:
     from mypy_boto3_secretsmanager import SecretsManagerClient
 
@@ -116,9 +118,49 @@ def _get_multiple(self, path: str, **sdk_options) -> Dict[str, str]:
         raise NotImplementedError()
 
 
+@overload
+def get_secret(
+    name: str,
+    transform: None = None,
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> str: ...
+
+
+@overload
+def get_secret(
+    name: str,
+    transform: Literal["json"],
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> dict: ...
+
+
+@overload
+def get_secret(
+    name: str,
+    transform: Literal["binary"],
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> Union[str, dict, bytes]: ...
+
+
+@overload
+def get_secret(
+    name: str,
+    transform: Literal["auto"],
+    force_fetch: bool = False,
+    max_age: Optional[int] = None,
+    **sdk_options,
+) -> bytes: ...
+
+
 def get_secret(
     name: str,
-    transform: Optional[str] = None,
+    transform: TransformOptions = None,
     force_fetch: bool = False,
     max_age: Optional[int] = None,
     **sdk_options,