fix: Add types for session_functions.py API call responses

Author: KShivendu

supertokens_python/recipe/session/recipe_implementation.py

@@ -37,7 +37,6 @@
     SessionDoesNotExistError,
     SessionInformationResult,
     SessionObj,
-    TokenInfo,
 )
 from .jwks import JWKClient
 from .jwt import ParsedJWTInfo, parse_jwt_without_signature_verification
@@ -93,27 +92,20 @@ async def create_new_session(
         log_debug_message("createNewSession: Finished")
 
         payload = parse_jwt_without_signature_verification(
-            result["accessToken"]["token"]
+            result.accessToken.token
         ).payload
 
-        refresh_token = result["refreshToken"]
-        refresh_token_info = TokenInfo(
-            refresh_token["token"],
-            refresh_token["expiry"],
-            refresh_token["createdTime"],
-        )
-
         new_session = Session(
             self,
             self.config,
-            result["accessToken"]["token"],
+            result.accessToken.token,
             build_front_token(
-                result["session"]["userId"], result["accessToken"]["expiry"], payload
+                result.session.userId, result.accessToken.expiry, payload
             ),
-            refresh_token_info,
-            result.get("antiCsrfToken"),
-            result["session"]["handle"],
-            result["session"]["userId"],
+            result.refreshToken,
+            result.antiCsrfToken,
+            result.session.handle,
+            result.session.userId,
             payload,
             None,
             True,
@@ -231,28 +223,28 @@ async def get_session(
 
         log_debug_message("getSession: Success!")
 
-        if "accessToken" in response:
+        if response.accessToken is not None:
             payload = parse_jwt_without_signature_verification(
-                response["accessToken"]["token"]
+                response.accessToken.token
             ).payload
-            access_token_str = response["accessToken"]["token"]
-            expiry_time = response["accessToken"]["expiry"]
+            access_token_str = response.accessToken.token
+            expiry_time = response.accessToken.expiry
             access_token_updated = True
         else:
             payload = access_token_obj.payload
             access_token_str = access_token
-            expiry_time = response["session"]["expiryTime"]
+            expiry_time = response.session.expiryTime
             access_token_updated = False
 
         session = Session(
             self,
             self.config,
             access_token_str,
-            build_front_token(response["session"]["userId"], expiry_time, payload),
+            build_front_token(response.session.userId, expiry_time, payload),
             None,  # refresh_token
             anti_csrf_token,
-            response["session"]["handle"],
-            response["session"]["userId"],
+            response.session.handle,
+            response.session.userId,
             payload,
             None,
             access_token_updated,
@@ -287,29 +279,22 @@ async def refresh_session(
         log_debug_message("refreshSession: Success!")
 
         payload = parse_jwt_without_signature_verification(
-            response["accessToken"]["token"]
+            response.accessToken.token,
         ).payload
 
-        new_refresh_token: Dict[str, Any] = response["refreshToken"]
-        new_refresh_token_info = TokenInfo(
-            new_refresh_token["token"],
-            new_refresh_token["expiry"],
-            new_refresh_token["createdTime"],
-        )
-
         session = Session(
             self,
             self.config,
-            response["accessToken"]["token"],
+            response.accessToken.token,
             build_front_token(
-                response["session"]["userId"],
-                response["accessToken"]["expiry"],
+                response.session.userId,
+                response.accessToken.expiry,
                 payload,
             ),
-            new_refresh_token_info,
-            response.get("antiCsrfToken"),
-            response["session"]["handle"],
-            response["session"]["userId"],
+            response.refreshToken,
+            response.antiCsrfToken,
+            response.session.handle,
+            response.session.userId,
             user_data_in_access_token=payload,
             req_res_info=None,
             access_token_updated=True,

supertokens_python/recipe/session/session_functions.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Union, Optional
 
 from supertokens_python.recipe.session.interfaces import SessionInformationResult
 
@@ -28,6 +28,7 @@
 from supertokens_python.logger import log_debug_message
 from supertokens_python.normalised_url_path import NormalisedURLPath
 from supertokens_python.process_state import AllowedProcessStates, ProcessState
+from supertokens_python.recipe.session.interfaces import TokenInfo
 
 from .exceptions import (
     TryRefreshTokenError,
@@ -37,13 +38,61 @@
 )
 
 
+class CreateOrRefreshAPIResponseSession:
+    def __init__(self, handle: str, userId: str, userDataInJWT: Any):
+        self.handle = handle
+        self.userId = userId
+        self.userDataInJWT = userDataInJWT
+
+
+class CreateOrRefreshAPIResponse:
+    def __init__(
+        self,
+        session: CreateOrRefreshAPIResponseSession,
+        accessToken: TokenInfo,
+        refreshToken: TokenInfo,
+        antiCsrfToken: Optional[str],
+    ):
+        self.session = session
+        self.accessToken = accessToken
+        self.refreshToken = refreshToken
+        self.antiCsrfToken = antiCsrfToken
+
+
+class GetSessionAPIResponseSession:
+    def __init__(
+        self, handle: str, userId: str, userDataInJWT: Dict[str, Any], expiryTime: int
+    ) -> None:
+        self.handle = handle
+        self.userId = userId
+        self.userDataInJWT = userDataInJWT
+        self.expiryTime = expiryTime
+
+
+class GetSessionAPIResponseAccessToken:
+    def __init__(self, token: str, expiry: int, createdTime: int) -> None:
+        self.token = token
+        self.expiry = expiry
+        self.createdTime = createdTime
+
+
+class GetSessionAPIResponse:
+    def __init__(
+        self,
+        session: GetSessionAPIResponseSession,
+        accessToken: Optional[GetSessionAPIResponseAccessToken] = None,
+    ) -> None:
+        self.session = session
+        self.accessToken = accessToken
+
+
 async def create_new_session(
     recipe_implementation: RecipeImplementation,
     user_id: str,
     disable_anti_csrf: bool,
     access_token_payload: Union[None, Dict[str, Any]],
     session_data_in_database: Union[None, Dict[str, Any]],
-) -> Dict[str, Any]:
+) -> CreateOrRefreshAPIResponse:
     if session_data_in_database is None:
         session_data_in_database = {}
     if access_token_payload is None:
@@ -66,7 +115,22 @@ async def create_new_session(
 
     response.pop("status", None)
 
-    return response  # FIXME: type the response
+    return CreateOrRefreshAPIResponse(
+        CreateOrRefreshAPIResponseSession(
+            response["handle"], response["userId"], response["userDataInJWT"]
+        ),
+        TokenInfo(
+            response["accessToken"]["token"],
+            response["accessToken"]["expiry"],
+            response["accessToken"]["createdTime"],
+        ),
+        TokenInfo(
+            response["refreshToken"]["token"],
+            response["refreshToken"]["expiry"],
+            response["refreshToken"]["createdTime"],
+        ),
+        response["antiCsrfToken"] if "antiCsrfToken" in response else None,
+    )
 
 
 async def get_session(
@@ -75,9 +139,9 @@ async def get_session(
     anti_csrf_token: Union[str, None],
     do_anti_csrf_check: bool,
     always_check_core: bool,
-) -> Dict[str, Any]:
+) -> GetSessionAPIResponse:
     config = recipe_implementation.config
-    access_token_info = None
+    access_token_info: Optional[Dict[str, Any]] = None
 
     try:
         access_token_info = get_info_from_access_token(
@@ -179,14 +243,14 @@ async def get_session(
         and not always_check_core
         and access_token_info["parentRefreshTokenHash1"] is None
     ):
-        return {
-            "session": {
-                "handle": access_token_info["sessionHandle"],
-                "userId": access_token_info["userId"],
-                "userDataInJWT": access_token_info["userData"],
-                "expiryTime": access_token_info["expiryTime"],
-            }
-        }
+        return GetSessionAPIResponse(
+            GetSessionAPIResponseSession(
+                access_token_info["sessionHandle"],
+                access_token_info["userId"],
+                access_token_info["userData"],
+                access_token_info["expiryTime"],
+            )
+        )
 
     ProcessState.get_instance().add_state(
         AllowedProcessStates.CALLING_SERVICE_IN_VERIFY
@@ -206,7 +270,19 @@ async def get_session(
     )
     if response["status"] == "OK":
         response.pop("status", None)
-        return response  # FIXME: type the response
+        return GetSessionAPIResponse(
+            GetSessionAPIResponseSession(
+                response["session"]["handle"],
+                response["session"]["userId"],
+                response["session"]["userData"],
+                response["session"]["expiresAt"],
+            ),
+            GetSessionAPIResponseAccessToken(
+                response["accessToken"]["token"],
+                response["accessToken"]["expiry"],
+                response["accessToken"]["createdTime"],
+            ),
+        )
     if response["status"] == "UNAUTHORISED":
         log_debug_message("getSession: Returning UNAUTHORISED because of core response")
         raise_unauthorised_exception(response["message"])
@@ -222,7 +298,7 @@ async def refresh_session(
     refresh_token: str,
     anti_csrf_token: Union[str, None],
     disable_anti_csrf: bool,
-) -> Dict[str, Any]:
+) -> CreateOrRefreshAPIResponse:
     data = {
         "refreshToken": refresh_token,
         "enableAntiCsrf": (
@@ -249,7 +325,22 @@ async def refresh_session(
     )
     if response["status"] == "OK":
         response.pop("status", None)
-        return response  # FIXME: type the response
+        return CreateOrRefreshAPIResponse(
+            CreateOrRefreshAPIResponseSession(
+                response["handle"], response["userId"], response["userDataInJWT"]
+            ),
+            TokenInfo(
+                response["accessToken"]["token"],
+                response["accessToken"]["expiry"],
+                response["accessToken"]["createdTime"],
+            ),
+            TokenInfo(
+                response["refreshToken"]["token"],
+                response["refreshToken"]["expiry"],
+                response["refreshToken"]["createdTime"],
+            ),
+            response["antiCsrfToken"] if "antiCsrfToken" in response else None,
+        )
     if response["status"] == "UNAUTHORISED":
         log_debug_message(
             "refreshSession: Returning UNAUTHORISED because of core response"