Spaces:
Running on Zero
Running on Zero
| """ | |
| Shared helpers for every feature: UI constants and token / login handling. | |
| Each feature (tab) receives `session_token` (a gr.State) from `app.py` and calls | |
| `resolve_token()` to pick a token in this order: session -> OAuth -> environment. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import re | |
| from dataclasses import dataclass | |
| from functools import lru_cache | |
| from typing import Any, Optional | |
| import gradio as gr | |
| from huggingface_hub import login, logout, whoami | |
| from core import hf_login | |
| # --------------------------------------------------------------------------- | |
| # Constants / UI configuration | |
| # --------------------------------------------------------------------------- | |
| #: Valid split-name regex of the `datasets` library (datasets.splits._split_re). | |
| SPLIT_RE = re.compile(r"^\w+(\.\w+)*$") | |
| CSS = """ | |
| .gradio-container {max-width: 1180px !important;} | |
| .run-btn {min-height: 52px; font-size: 1.05rem; font-weight: 600;} | |
| .log-box textarea {font-family: ui-monospace, SFMono-Regular, Menlo, monospace; font-size: .82rem;} | |
| """ | |
| GRADIO_MAJOR = int(str(getattr(gr, "__version__", "5")).split(".")[0]) | |
| #: A job's status markdown starts with this when it failed (see `features/jobs.py`). | |
| FAILED_PREFIX = "### ❌" | |
| #: Gradio 6 renamed `show_copy_button` -> `buttons`; stay compatible with Gradio 5 and 6. | |
| COPY_BUTTON_KWARGS = {"buttons": ["copy"]} if GRADIO_MAJOR >= 6 else {"show_copy_button": True} | |
| # --------------------------------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------------------------------- | |
| def _clean(value: Optional[str]) -> str: | |
| """Normalize a string input.""" | |
| return value.strip() if isinstance(value, str) else "" | |
| def _as_token(value) -> Optional[str]: | |
| """Accept both `str` and `gr.OAuthToken` -> return the token string.""" | |
| if value is None: | |
| return None | |
| if isinstance(value, str): | |
| return value.strip() or None | |
| token = getattr(value, "token", None) # gr.OAuthToken | |
| return str(token).strip() if token else None | |
| def resolve_token(*candidates) -> Optional[str]: | |
| """Pick a token by priority: session -> OAuth -> Space environment variables.""" | |
| for candidate in list(candidates) + [ | |
| os.environ.get("HF_TOKEN"), | |
| os.environ.get("HUGGING_FACE_HUB_TOKEN"), | |
| ]: | |
| token = _as_token(candidate) | |
| if token: | |
| return token | |
| return None | |
| def token_user(token: Optional[str]) -> Optional[str]: | |
| """Return the username if the token is valid, otherwise None.""" | |
| if not token: | |
| return None | |
| try: | |
| info = whoami(token=token) | |
| except Exception: # noqa: BLE001 | |
| return None | |
| if isinstance(info, dict): | |
| return info.get("name") or info.get("user") or "unknown" | |
| return "unknown" | |
| def cached_user(token: Optional[str]) -> Optional[str]: | |
| """`token_user` cached per token — for handlers that run often (e.g. the Jobs tab timer).""" | |
| return token_user(token) | |
| def profile_username(profile) -> Optional[str]: | |
| """Get the username from a `gr.OAuthProfile` (dict-like).""" | |
| if not profile: | |
| return None | |
| for key in ("preferred_username", "username", "name", "sub"): | |
| try: | |
| value = profile.get(key) if hasattr(profile, "get") else getattr(profile, key, None) | |
| except Exception: # noqa: BLE001 | |
| value = None | |
| if value: | |
| return str(value) | |
| return None | |
| class AppContext: | |
| """Components shared by every tab, passed to each feature's `build(ctx)`.""" | |
| #: Token of this browser session (gr.State); resolve it with `resolve_token`. | |
| session_token: Any | |
| #: gr.State set to a job id whenever a tab submits a job (the Jobs tab follows it). | |
| submitted_job: Any | |
| #: The gr.Tabs container, so a handler can switch tabs. | |
| tabs: Any = None | |
| def login_badge() -> str: | |
| return "⚠️ **Not signed in.** Use the OAuth button above, or paste an Access Token." | |
| # --------------------------------------------------------------------------- | |
| # Sign in / sign out | |
| # --------------------------------------------------------------------------- | |
| def on_page_load( | |
| profile: Optional[gr.OAuthProfile], | |
| oauth_token: Optional[gr.OAuthToken], | |
| session_token: Optional[str], | |
| request: gr.Request, | |
| ): | |
| """Runs on page load; also finishes the OAuth flow when this page is its redirect.""" | |
| token, error = hf_login.handle_callback(request) | |
| if error: | |
| return session_token, f"❌ **Sign-in with Hugging Face failed:** {error}" | |
| if token: | |
| user = token_user(token) | |
| return token, ( | |
| f"✅ **Signed in:** `{user}` (Hugging Face OAuth). The page where you clicked " | |
| "*Sign in* is signed in as well — you can close this tab or keep using it." | |
| ) | |
| token = resolve_token(session_token, oauth_token) | |
| user = token_user(token) | |
| if user: | |
| scope = getattr(oauth_token, "scope", "") or "" | |
| extra = f" _(OAuth scope: `{scope}`)_" if scope else "" | |
| return session_token, f"✅ **Signed in:** `{user}`{extra}" | |
| name = profile_username(profile) | |
| if name: | |
| return session_token, ( | |
| f"👤 OAuth account: `{name}` — but no access token with write permission was received. " | |
| "Grant the `write-repos` / `manage-repos` scopes, or paste an Access Token." | |
| ) | |
| return session_token, login_badge() | |
| def start_oauth_login(nonce: Optional[str]): | |
| """After the sign-in button opened the Hugging Face tab: start polling for the token.""" | |
| if not nonce: | |
| return gr.skip(), gr.skip() | |
| if not hf_login.ENABLED: | |
| return ( | |
| "❌ Sign-in with Hugging Face only works on a Space with `hf_oauth: true` — " | |
| "paste an Access Token instead.", | |
| gr.Timer(active=False), | |
| ) | |
| hf_login.start(nonce) | |
| return ( | |
| "⏳ Waiting for the Hugging Face sign-in in the new tab (allow pop-ups if no tab opened)…", | |
| gr.Timer(active=True), | |
| ) | |
| def poll_oauth_login(nonce: Optional[str], session_token: Optional[str]): | |
| """Timer tick: claim the token parked by the OAuth redirect tab, if it arrived.""" | |
| stop = gr.Timer(active=False) | |
| if not nonce: | |
| return gr.skip(), gr.skip(), gr.skip(), stop | |
| token = hf_login.claim(nonce) | |
| if token: | |
| user = token_user(token) | |
| return token, f"✅ **Signed in:** `{user}` (Hugging Face OAuth)", "", stop | |
| if hf_login.expired(nonce): | |
| return gr.skip(), "⌛ Sign-in with Hugging Face timed out — click the button again.", "", stop | |
| return gr.skip(), gr.skip(), gr.skip(), gr.skip() | |
| def do_token_login(raw_token: str, session_token: Optional[str]): | |
| """Sign in with a manually pasted Access Token.""" | |
| token = _clean(raw_token) or _as_token(session_token) | |
| if not token: | |
| return session_token, "❌ No token entered." | |
| user = token_user(token) | |
| if not user: | |
| return ( | |
| session_token, | |
| "❌ Invalid or expired token. Create a new one at " | |
| "https://huggingface.co/settings/tokens (needs **Write** permission).", | |
| ) | |
| try: # best-effort: lets underlying libraries reuse it when no token is passed | |
| login(token=token, add_to_git_credential=False) | |
| except Exception: # noqa: BLE001 | |
| pass | |
| return token, f"✅ **Signed in:** `{user}`" | |
| def do_logout(session_token: Optional[str]): | |
| if session_token: | |
| try: | |
| logout(token=session_token) | |
| except Exception: # noqa: BLE001 | |
| pass | |
| return None, "👋 Token removed from this session." | |