From 611efdcd66e1bb6f1a05892fd2fd84c0f2b295c5 Mon Sep 17 00:00:00 2001 From: Thomas Marwitz Date: Tue, 10 Dec 2024 13:07:44 +0100 Subject: [PATCH] Type hinted credential dict --- minimalkv/net/_aws_refreshable_session.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/minimalkv/net/_aws_refreshable_session.py b/minimalkv/net/_aws_refreshable_session.py index 41d6b46..cb65d7e 100644 --- a/minimalkv/net/_aws_refreshable_session.py +++ b/minimalkv/net/_aws_refreshable_session.py @@ -3,7 +3,7 @@ from datetime import datetime from logging import getLogger -from typing import Dict, Optional +from typing import Any, Callable, Coroutine, Dict, Optional, TypedDict from aiobotocore.credentials import ( AioRefreshableCredentials, @@ -17,6 +17,13 @@ logger = getLogger(__name__) +class BotoCredentials(TypedDict): + access_key: str + secret_key: str + token: str + expiry_time: str + + class RefreshableAssumeRoleProvider(CredentialProvider): """A credential provider that loads self-refreshing credentials using 'sts.assume_role'. @@ -58,13 +65,13 @@ async def load(self) -> AioRefreshableCredentials: return refreshable_credentials - def create_refresh(self): - async def _refresh() -> Dict: + def create_refresh(self) -> Callable[[], Coroutine[Any, Any, BotoCredentials]]: + async def _refresh() -> BotoCredentials: logger.info("Refreshing credentials") sts_client = self.sts_session.create_client("sts") refresh = create_assume_role_refresher(sts_client, self._assume_role_params) - credentials = await refresh() + credentials: BotoCredentials = await refresh() to_zone = tz.tzlocal() local_expiry_time = ( datetime.strptime(credentials["expiry_time"], "%Y-%m-%dT%H:%M:%S%Z")