add typed dicts for convenience

Author: grdsdev

supabase_auth/_async/gotrue_client.py

@@ -43,6 +43,7 @@
 from ..http_clients import AsyncClient
 from ..timer import Timer
 from ..types import (
+    JWK,
     AuthChangeEvent,
     AuthenticatorAssuranceLevels,
     AuthFlowType,
@@ -58,6 +59,7 @@
     CodeExchangeParams,
     DecodedJWTDict,
     IdentitiesResponse,
+    JWKSet,
     MFAChallengeAndVerifyParams,
     MFAChallengeParams,
     MFAEnrollParams,
@@ -111,7 +113,7 @@ def __init__(
             verify=verify,
             proxy=proxy,
         )
-        self._jwks = {"keys": []}
+        self._jwks: JWKSet = {"keys": []}
         self._storage_key = storage_key or STORAGE_KEY
         self._auto_refresh_token = auto_refresh_token
         self._persist_session = persist_session
@@ -1158,19 +1160,19 @@ async def exchange_code_for_session(self, params: CodeExchangeParams):
             self._notify_all_subscribers("SIGNED_IN", response.session)
         return response
 
-    async def _fetch_jwks(self, kid: str, jwks: Dict[str, list]) -> Dict[str, Any]:
-        jwk: Dict[str, Any] = {}
+    async def _fetch_jwks(self, kid: str, jwks: JWKSet) -> JWK:
+        jwk: Optional[JWK] = None
 
         # try fetching from the suplied keys.
-        jwk = next((jwk for jwk in jwks.get("keys", []) if jwk.get("kid") == kid), None)
+        jwk = next((jwk for jwk in jwks["keys"] if jwk["kid"] == kid), None)
 
         if jwk:
             return jwk
 
         if self._jwks:
             # try fetching from the cache.
             jwk = next(
-                (jwk for jwk in self._jwks.get("keys", []) if jwk.get("kid") == kid),
+                (jwk for jwk in self._jwks["keys"] if jwk["kid"] == kid),
                 None,
             )
             if jwk:
@@ -1182,9 +1184,7 @@ async def _fetch_jwks(self, kid: str, jwks: Dict[str, list]) -> Dict[str, Any]:
             self._jwks = response
 
             # find the signing key
-            jwk = next(
-                (jwk for jwk in response.get("keys", []) if jwk.get("kid") == kid), None
-            )
+            jwk = next((jwk for jwk in response["keys"] if jwk["kid"] == kid), None)
             if not jwk:
                 raise AuthInvalidJwtError("No matching signing key found in JWKS")
 
@@ -1193,7 +1193,7 @@ async def _fetch_jwks(self, kid: str, jwks: Dict[str, list]) -> Dict[str, Any]:
         raise AuthInvalidJwtError("JWT has no valid kid")
 
     async def get_claims(
-        self, jwt: Optional[str] = None, jwks: Optional[Dict[str, list]] = None
+        self, jwt: Optional[str] = None, jwks: Optional[JWKSet] = None
     ) -> Optional[ClaimsResponse]:
         token = jwt
         if not token:
@@ -1204,12 +1204,16 @@ async def get_claims(
             token = session.access_token
 
         decoded_jwt = decode_jwt(token)
-        payload = decoded_jwt["payload"]
-        header = decoded_jwt["header"]
-        signature = decoded_jwt["signature"]
 
-        raw_header = decoded_jwt["raw"]["header"]
-        raw_payload = decoded_jwt["raw"]["payload"]
+        payload, header, signature = (
+            decoded_jwt["payload"],
+            decoded_jwt["header"],
+            decoded_jwt["signature"],
+        )
+        raw_header, raw_payload = (
+            decoded_jwt["raw"]["header"],
+            decoded_jwt["raw"]["payload"],
+        )
 
         validate_exp(payload["exp"])
 
@@ -1220,7 +1224,7 @@ async def get_claims(
 
         algorithm = get_algorithm_by_name(header["alg"])
         signing_key = algorithm.from_jwk(
-            await self._fetch_jwks(header["kid"], jwks or {})
+            await self._fetch_jwks(header["kid"], jwks or {"keys": []})
         )
 
         # verify the signature
@@ -1235,8 +1239,8 @@ async def get_claims(
         return ClaimsResponse(claims=payload, headers=header, signature=signature)
 
 
-def parse_jwks(response: Any) -> Dict[str, list]:
+def parse_jwks(response: Any) -> JWKSet:
     if "keys" not in response or len(response["keys"]) == 0:
         raise AuthInvalidJwtError("JWKS is empty")
 
-    return response
+    return {"keys": response["keys"]}

supabase_auth/_sync/gotrue_client.py

@@ -43,6 +43,7 @@
 from ..http_clients import SyncClient
 from ..timer import Timer
 from ..types import (
+    JWK,
     AuthChangeEvent,
     AuthenticatorAssuranceLevels,
     AuthFlowType,
@@ -58,6 +59,7 @@
     CodeExchangeParams,
     DecodedJWTDict,
     IdentitiesResponse,
+    JWKSet,
     MFAChallengeAndVerifyParams,
     MFAChallengeParams,
     MFAEnrollParams,
@@ -111,7 +113,7 @@ def __init__(
             verify=verify,
             proxy=proxy,
         )
-        self._jwks = {"keys": []}
+        self._jwks: JWKSet = {"keys": []}
         self._storage_key = storage_key or STORAGE_KEY
         self._auto_refresh_token = auto_refresh_token
         self._persist_session = persist_session
@@ -421,7 +423,9 @@ def sign_in_with_oauth(
         )
         return OAuthResponse(provider=provider, url=url_with_qs)
 
-    def link_identity(self, credentials: SignInWithOAuthCredentials) -> OAuthResponse:
+    def link_identity(
+        self, credentials: SignInWithOAuthCredentials
+    ) -> OAuthResponse:
         provider = credentials.get("provider")
         options = credentials.get("options", {})
         redirect_to = options.get("redirect_to")
@@ -704,7 +708,9 @@ def set_session(self, access_token: str, refresh_token: str) -> AuthResponse:
         self._notify_all_subscribers("TOKEN_REFRESHED", session)
         return AuthResponse(session=session, user=response.user)
 
-    def refresh_session(self, refresh_token: Optional[str] = None) -> AuthResponse:
+    def refresh_session(
+        self, refresh_token: Optional[str] = None
+    ) -> AuthResponse:
         """
         Returns a new session, regardless of expiry status.
 
@@ -1113,7 +1119,9 @@ def _get_url_for_provider(
         if self._flow_type == "pkce":
             code_verifier = generate_pkce_verifier()
             code_challenge = generate_pkce_challenge(code_verifier)
-            self._storage.set_item(f"{self._storage_key}-code-verifier", code_verifier)
+            self._storage.set_item(
+                f"{self._storage_key}-code-verifier", code_verifier
+            )
             code_challenge_method = (
                 "plain" if code_verifier == code_challenge else "s256"
             )
@@ -1152,19 +1160,19 @@ def exchange_code_for_session(self, params: CodeExchangeParams):
             self._notify_all_subscribers("SIGNED_IN", response.session)
         return response
 
-    def _fetch_jwks(self, kid: str, jwks: Dict[str, list]) -> Dict[str, Any]:
-        jwk: Dict[str, Any] = {}
+    def _fetch_jwks(self, kid: str, jwks: JWKSet) -> JWK:
+        jwk: Optional[JWK] = None
 
         # try fetching from the suplied keys.
-        jwk = next((jwk for jwk in jwks.get("keys", []) if jwk.get("kid") == kid), None)
+        jwk = next((jwk for jwk in jwks["keys"] if jwk["kid"] == kid), None)
 
         if jwk:
             return jwk
 
         if self._jwks:
             # try fetching from the cache.
             jwk = next(
-                (jwk for jwk in self._jwks.get("keys", []) if jwk.get("kid") == kid),
+                (jwk for jwk in self._jwks["keys"] if jwk["kid"] == kid),
                 None,
             )
             if jwk:
@@ -1176,9 +1184,7 @@ def _fetch_jwks(self, kid: str, jwks: Dict[str, list]) -> Dict[str, Any]:
             self._jwks = response
 
             # find the signing key
-            jwk = next(
-                (jwk for jwk in response.get("keys", []) if jwk.get("kid") == kid), None
-            )
+            jwk = next((jwk for jwk in response["keys"] if jwk["kid"] == kid), None)
             if not jwk:
                 raise AuthInvalidJwtError("No matching signing key found in JWKS")
 
@@ -1187,7 +1193,7 @@ def _fetch_jwks(self, kid: str, jwks: Dict[str, list]) -> Dict[str, Any]:
         raise AuthInvalidJwtError("JWT has no valid kid")
 
     def get_claims(
-        self, jwt: Optional[str] = None, jwks: Optional[Dict[str, list]] = None
+        self, jwt: Optional[str] = None, jwks: Optional[JWKSet] = None
     ) -> Optional[ClaimsResponse]:
         token = jwt
         if not token:
@@ -1198,12 +1204,16 @@ def get_claims(
             token = session.access_token
 
         decoded_jwt = decode_jwt(token)
-        payload = decoded_jwt["payload"]
-        header = decoded_jwt["header"]
-        signature = decoded_jwt["signature"]
 
-        raw_header = decoded_jwt["raw"]["header"]
-        raw_payload = decoded_jwt["raw"]["payload"]
+        payload, header, signature = (
+            decoded_jwt["payload"],
+            decoded_jwt["header"],
+            decoded_jwt["signature"],
+        )
+        raw_header, raw_payload = (
+            decoded_jwt["raw"]["header"],
+            decoded_jwt["raw"]["payload"],
+        )
 
         validate_exp(payload["exp"])
 
@@ -1213,7 +1223,9 @@ def get_claims(
             return ClaimsResponse(claims=payload, headers=header, signature=signature)
 
         algorithm = get_algorithm_by_name(header["alg"])
-        signing_key = algorithm.from_jwk(self._fetch_jwks(header["kid"], jwks or {}))
+        signing_key = algorithm.from_jwk(
+            self._fetch_jwks(header["kid"], jwks or {"keys": []})
+        )
 
         # verify the signature
         is_valid = algorithm.verify(
@@ -1227,8 +1239,8 @@ def get_claims(
         return ClaimsResponse(claims=payload, headers=header, signature=signature)
 
 
-def parse_jwks(response: Any) -> Dict[str, list]:
+def parse_jwks(response: Any) -> JWKSet:
     if "keys" not in response or len(response["keys"]) == 0:
         raise AuthInvalidJwtError("JWKS is empty")
 
-    return response
+    return {"keys": response["keys"]}

supabase_auth/types.py

@@ -816,6 +816,17 @@ class ClaimsResponse(TypedDict):
     signature: bytes
 
 
+class JWK(TypedDict, total=False):
+    kty: Literal["RSA", "EC", "oct"]
+    key_ops: List[str]
+    alg: Optional[str]
+    kid: Optional[str]
+
+
+class JWKSet(TypedDict):
+    keys: List[JWK]
+
+
 for model in [
     AMREntry,
     AuthResponse,

tests/_sync/test_gotrue.py

@@ -1,41 +1,38 @@
 import unittest
-import pytest
-from pytest_mock import mocker
 
+from .clients import auth_client, auth_client_with_asymmetric_session
 from .utils import mock_user_credentials
 
-from .clients import auth_client, auth_client_with_asymmetric_session
 
 def test_get_claims_returns_none_when_session_is_none():
-  claims = auth_client().get_claims()
-  assert claims is None
+    claims = auth_client().get_claims()
+    assert claims is None
+
 
 def test_get_claims_calls_get_user_if_symmetric_jwt(mocker):
-  client = auth_client()
-  spy = mocker.spy(client, 'get_user')
+    client = auth_client()
+    spy = mocker.spy(client, "get_user")
 
-  user = (client.sign_up(mock_user_credentials())).user
-  assert user is not None
+    user = (client.sign_up(mock_user_credentials())).user
+    assert user is not None
 
-  claims = (client.get_claims())["claims"]
-  assert claims["email"] == user.email
-  spy.assert_called_once()
-  
+    claims = (client.get_claims())["claims"]
+    assert claims["email"] == user.email
+    spy.assert_called_once()
 
-def test_get_claims_fetches_jwks_to_verify_asymmetric_jwt(mocker):
-  client = auth_client_with_asymmetric_session()
 
-  user = (client.sign_up(mock_user_credentials())).user
-  assert user is not None
+def test_get_claims_fetches_jwks_to_verify_asymmetric_jwt(mocker):
+    client = auth_client_with_asymmetric_session()
 
-  spy = mocker.spy(client, "_request")
+    user = (client.sign_up(mock_user_credentials())).user
+    assert user is not None
 
-  claims = (client.get_claims())["claims"]
-  assert claims["email"] == user.email
+    spy = mocker.spy(client, "_request")
 
-  spy.assert_called_once()
-  spy.assert_called_with("GET", ".well-known/jwks.json", xform=unittest.mock.ANY)
+    claims = (client.get_claims())["claims"]
+    assert claims["email"] == user.email
 
-  assert len(spy.spy_return.get("keys")) > 0
+    spy.assert_called_once()
+    spy.assert_called_with("GET", ".well-known/jwks.json", xform=unittest.mock.ANY)
 
-  
+    assert len(spy.spy_return.get("keys")) > 0

tests/_sync/test_gotrue_admin_api.py

@@ -156,27 +156,23 @@ def test_modify_confirm_email_using_update_user_by_id():
 
 def test_invalid_credential_sign_in_with_phone():
     try:
-        response = (
-            client_api_auto_confirm_off_signups_enabled_client().sign_in_with_password(
-                {
-                    "phone": "+123456789",
-                    "password": "strong_pwd",
-                }
-            )
+        response = client_api_auto_confirm_off_signups_enabled_client().sign_in_with_password(
+            {
+                "phone": "+123456789",
+                "password": "strong_pwd",
+            }
         )
     except AuthApiError as e:
         assert e.to_dict()
 
 
 def test_invalid_credential_sign_in_with_email():
     try:
-        response = (
-            client_api_auto_confirm_off_signups_enabled_client().sign_in_with_password(
-                {
-                    "email": "[email protected]",
-                    "password": "strong_pwd",
-                }
-            )
+        response = client_api_auto_confirm_off_signups_enabled_client().sign_in_with_password(
+            {
+                "email": "[email protected]",
+                "password": "strong_pwd",
+            }
         )
     except AuthApiError as e:
         assert e.to_dict()
@@ -390,10 +386,12 @@ def test_sign_in_with_sso():
 
 
 def test_sign_in_with_oauth():
-    assert client_api_auto_confirm_off_signups_enabled_client().sign_in_with_oauth(
-        {
-            "provider": "google",
-        }
+    assert (
+        client_api_auto_confirm_off_signups_enabled_client().sign_in_with_oauth(
+            {
+                "provider": "google",
+            }
+        )
     )