cleanup: pydantic for parsing claims

Author: mrkaye97

hatchet_sdk/token.py

@@ -1,28 +1,31 @@
 import base64
-import json
-from typing import Any
 
+from pydantic import BaseModel
 
-def get_tenant_id_from_jwt(token: str) -> str:
-    claims = extract_claims_from_jwt(token)
 
-    return claims.get("sub")
+class Claims(BaseModel):
+    sub: str
+    server_url: str
+    grpc_broadcast_address: str
+
+
+def get_tenant_id_from_jwt(token: str) -> str:
+    return extract_claims_from_jwt(token).sub
 
 
 def get_addresses_from_jwt(token: str) -> tuple[str, str]:
     claims = extract_claims_from_jwt(token)
 
-    return claims["server_url"], claims["grpc_broadcast_address"]
+    return claims.server_url, claims.grpc_broadcast_address
 
 
-def extract_claims_from_jwt(token: str) -> dict[str, Any]:
+def extract_claims_from_jwt(token: str) -> Claims:
     parts = token.split(".")
     if len(parts) != 3:
         raise ValueError("Invalid token format")
 
     claims_part = parts[1]
     claims_part += "=" * ((4 - len(claims_part) % 4) % 4)  # Padding for base64 decoding
     claims_data = base64.urlsafe_b64decode(claims_part)
-    claims = json.loads(claims_data)
 
-    return claims
+    return Claims.model_validate_json(claims_data)