feat(auth): add support for Google Colab token retrieval (#377)

* feat(auth): add support for Google Colab token retrieval

* feat(auth): implement Google Colab token retrieval mechanism

* fix(auth): correct conditional logic in fetch_credentials function

* fix(auth): remove unnecessary noqa comment in is_google_colab function

* fix(auth): remove unnecessary noqa comment and import get_ipython

* fix(auth): handle token parsing and improve Google Colab detection

* fix(auth): correct conditional logic in fetch_credentials function

* fix(auth): reorder credential fetching logic to prioritize Google Colab token

* fix(auth): remove Google Colab token handling and related logic from fal

* fix(auth): remove Google Colab state management and token retrieval logic

* fix(auth): implement Google Colab token retrieval with state management

projects/fal_client/src/fal_client/auth.py

@@ -1,4 +1,49 @@
 import os
+from threading import Lock
+from typing import Optional
+
+
+class GoogleColabState:
+    def __init__(self):
+        self.is_checked = False
+        self.lock = Lock()
+        self.secret: Optional[str] = None
+
+
+_colab_state = GoogleColabState()
+
+
+def is_google_colab() -> bool:
+    try:
+        from IPython import get_ipython
+
+        return "google.colab" in str(get_ipython())
+    except ModuleNotFoundError:
+        return False
+    except NameError:
+        return False
+
+
+def get_colab_token() -> Optional[str]:
+    if not is_google_colab():
+        return None
+    with _colab_state.lock:
+        if _colab_state.is_checked:  # request access only once
+            return _colab_state.secret
+
+        try:
+            from google.colab import userdata  # noqa: I001
+        except ImportError:
+            return None
+
+        try:
+            token = userdata.get("FAL_KEY")
+            _colab_state.secret = token.strip()
+        except Exception:
+            _colab_state.secret = None
+
+        _colab_state.is_checked = True
+        return _colab_state.secret
 
 
 class MissingCredentialsError(Exception):
@@ -15,6 +60,8 @@ def fetch_credentials() -> str:
         fal_key_secret := os.getenv("FAL_KEY_SECRET")
     ):
         return f"{key_id}:{fal_key_secret}"
+    elif colab_token := get_colab_token():
+        return colab_token
     else:
         raise MissingCredentialsError(
             "Please set the FAL_KEY environment variable to your API key, or set the FAL_KEY_ID and FAL_KEY_SECRET environment variables."