From 42dd942b867936a95604701a6717fbdbdfd63f94 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Mon, 20 Jan 2025 20:10:37 +0200
Subject: [PATCH 1/2] feat(fal/auth): add support for Google Colab token
 retrieval

Adapting #377 to work with sdk as well.

Kudos @not-lain

Co-authored-by: Hafedh Hichri <70411813+not-lain@users.noreply.github.com>
---
 projects/fal/src/fal/auth/__init__.py | 49 +++++++++++++++++++++++++--
 1 file changed, 47 insertions(+), 2 deletions(-)

diff --git a/projects/fal/src/fal/auth/__init__.py b/projects/fal/src/fal/auth/__init__.py
index de83a6e1..91e3f662 100644
--- a/projects/fal/src/fal/auth/__init__.py
+++ b/projects/fal/src/fal/auth/__init__.py
@@ -2,6 +2,8 @@
 
 import os
 from dataclasses import dataclass, field
+from threading import Lock
+from typing import Optional
 
 import click
 
@@ -11,13 +13,56 @@
 from fal.exceptions.auth import UnauthenticatedException
 
 
+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
+
+
 def key_credentials() -> tuple[str, str] | None:
     # Ignore key credentials when the user forces auth by user.
     if os.environ.get("FAL_FORCE_AUTH_BY_USER") == "1":
         return None
 
-    if "FAL_KEY" in os.environ:
-        key = os.environ["FAL_KEY"]
+    key = os.environ.get("FAL_KEY") or get_colab_token()
+    if key:
         key_id, key_secret = key.split(":", 1)
         return (key_id, key_secret)
     elif "FAL_KEY_ID" in os.environ and "FAL_KEY_SECRET" in os.environ:

From 8d992f130981695094319275bdf03262933c84d8 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Tue, 21 Jan 2025 15:39:44 +0200
Subject: [PATCH 2/2] fix(fal): use importlib.metadat instead of deprecated
 pkg_resources

---
 projects/fal/src/fal/app.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/projects/fal/src/fal/app.py b/projects/fal/src/fal/app.py
index 9cb197ee..86c1ca89 100644
--- a/projects/fal/src/fal/app.py
+++ b/projects/fal/src/fal/app.py
@@ -249,9 +249,9 @@ def _to_fal_app_name(name: str) -> str:
 
 
 def _print_python_packages() -> None:
-    from pkg_resources import working_set
+    from importlib.metadata import distributions
 
-    packages = [f"{package.key}=={package.version}" for package in working_set]
+    packages = [f"{dist.metadata['Name']}=={dist.version}" for dist in distributions()]
 
     print("[debug] Python packages installed:", ", ".join(packages))