""" 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" @lru_cache(maxsize=256) 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 @dataclass 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."