Add AuthProvider construction to MetafoldClient

Author: rhamel3d

metafold/__init__.py

@@ -2,6 +2,7 @@
 from metafold.projects import ProjectsEndpoint
 from metafold.assets import AssetsEndpoint
 from metafold.jobs import JobsEndpoint
+from metafold.auth import AuthProvider
 from typing import Optional
 
 
@@ -18,8 +19,12 @@ class MetafoldClient(Client):
     jobs: JobsEndpoint
 
     def __init__(
-        self, access_token: str,
+        self,
+        access_token: Optional[str] = None,
         project_id: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        auth_domain: str = "metafold3d.us.auth0.com",
         base_url: str = "https://api.metafold3d.com",
     ) -> None:
         """Initialize Metafold API client.
@@ -29,7 +34,17 @@ def __init__(
             project_id: ID of the project to make API calls against.
             base_url: Metafold API URL. Used for internal testing.
         """
-        super().__init__(access_token, base_url, project_id=project_id)
+        # client_id and client_secret have priority
+        if not any([client_id and client_secret, access_token]):
+            raise ValueError(
+                "Expected client_id and client_secret or access_token to be provided"
+            )
+        elif client_id and client_secret:
+            auth = AuthProvider(client_id, client_secret, auth_domain, base_url)
+            super().__init__(base_url, auth=auth, project_id=project_id)
+        else:
+            super().__init__(base_url, access_token=access_token, project_id=project_id)
+
         self.projects = ProjectsEndpoint(self)
         self.assets = AssetsEndpoint(self)
         self.jobs = JobsEndpoint(self)

metafold/auth.py

@@ -1,13 +1,21 @@
 from auth0.authentication import GetToken
 from collections import namedtuple
 from datetime import datetime, timedelta, timezone
+from typing import Optional
 
 Token = namedtuple("Token", ["access_token", "expires_at"])
 
 
 class AuthProvider:
-    def __init__(self, auth_domain: str, client_id: str, client_secret: str):
+    def __init__(
+        self,
+        client_id: str,
+        client_secret: str,
+        auth_domain: str,
+        base_url: str
+    ) ->  None:
         self._auth_domain = auth_domain
+        self._base_url = base_url
         self._client_id = client_id
 
         self._get_token = GetToken(
@@ -20,7 +28,7 @@ def __init__(self, auth_domain: str, client_id: str, client_secret: str):
     def get_token(self) -> str:
         now = datetime.now(timezone.utc)
         if not self._token or self._token.expires_at - now < timedelta(minutes=1):
-            token = get_token.client_credentials(self._base_url)
+            token = self._get_token.client_credentials(self._base_url)
             expires_at = now + timedelta(seconds=token["expires_in"])
             self._token = Token(token["access_token"], expires_at)
         return self._token.access_token

metafold/client.py

@@ -1,3 +1,4 @@
+from metafold.auth import AuthProvider
 from requests import HTTPError, Response, Session
 from typing import Any, Callable, Optional
 from urllib.parse import urljoin
@@ -8,17 +9,29 @@ class Client:
     """Base client."""
 
     def __init__(
-        self, access_token: str, base_url: str,
+        self,
+        base_url: str,
+        access_token: Optional[str] = None,
         project_id: Optional[str] = None,
+        auth: Optional[AuthProvider] = None
     ) -> None:
+        if bool(auth) == bool(access_token):
+            raise ValueError(
+                "Expected AuthProvider or access_token to be provided"
+            )
+        self._auth = auth
         self._default_project = project_id
         self._base_url = base_url
         self._session = Session()
         self._session.headers.update({
             "Accept": "application/json",
-            "Authorization": f"Bearer {access_token}",
             "User-Agent": f"Python/{platform.python_version()}",
         })
+        if access_token:
+            self._session.headers.update({
+                "Authorization": f"Bearer {access_token}",
+            })
+
 
     def project_id(self, id: Optional[str] = None) -> str:
         id = id or self._default_project
@@ -33,7 +46,10 @@ def _request(
         *args: Any, **kwargs: Any,
     ) -> Response:
         url = urljoin(self._base_url, url)
-        r: Response = request(url, *args, **kwargs)
+        headers = None
+        if self._auth:
+            headers = {"Authorization": f"Bearer {self._auth.get_token()}"}
+        r: Response = request(url, *args, **kwargs, headers=headers)
         if not r.ok:
             body: dict[str, Any] = r.json()
             # Error responses aren't entirely consistent in the Metafold API,