open_pr.py

import re
import time

import gradio as gr
from huggingface_hub import CommitOperationAdd, HfApi, RepoUrl, SpaceCard

IMPORTS_REGEX = re.compile(r"^(from|import) .*$", re.MULTILINE)

PR_DESCRIPTION = """
This PR enables Space CI on your Space. **Gradio Space CI is a tool to create ephemeral Spaces for each PR opened on your Space repo.** The goal is to improve developer experience by making the review process as lean as possible.

### ⚙️ How it works:

- Listens to pull requests events:
  - If PR is opened => starts an ephemeral Space
  - If PR is updated => updates the Space
  - If PR is closed => deleted the Space
- Checks PR author:
  - If trusted author => ephemeral Space is configured with variables, secrets and hardware.
  - If not a trusted author => ephemeral Space is started without configuration.
  - Space owners are trusted by default. Additional "trusted authors" can be configuration manually.

### ⚠️ Before merging:

1. Check that the configuration is correct. By default the Space is configured to run ephemeral Spaces on a (free) CPU instance without any secrets.
2. You must set `HF_TOKEN` as a secret in your Space settings. Token must have 'write' permission. You can create a new one in your [User settings](https://huggingface.co/settings/token).

---
This is an automated PR created with https://huggingface.co/spaces/Wauplin/gradio-space-ci.
For more details about Space CI, checkout [this page]](https://huggingface.co/spaces/Wauplin/gradio-space-ci/blob/main/README.md).
If you find any issues, please report here: https://huggingface.co/spaces/Wauplin/gradio-space-ci/discussions

Feel free to ignore this PR.
"""

SUCCESS_MESSAGE = """
### Success 🔥

Yay! A PR has been open to enable Space CI on {space_id}. Check it out here: [{pr_url}]({pr_url}).

You can contact the Space owner to let them know about this PR.
"""


def open_pr(space_id_or_url: str, oauth_token: gr.OAuthToken | None) -> str:
    if oauth_token is None:
        raise gr.Error("You must be logged in to open a PR. Click on 'Sign in with Huggingface' first.")
    token = oauth_token.token
    if oauth_token.expires_at > time.time():
        raise gr.Error("Token expired. Logout and try again.")

    api = HfApi(token=token)

    if not space_id:
        raise gr.Error("You must input a Space ID or URL.")
    space_id = (
        RepoUrl(space_id_or_url).repo_id if "https://huggingface.co/spaces/" in space_id_or_url else space_id_or_url
    )

    # 0. Check token + repo existence
    try:
        user = api.whoami()
        if user["type"] != "user":
            raise gr.Error(
                "You must use a user token, not an organization token. Go to https://huggingface.co/settings/token to create a new token."
            )
    except Exception as e:
        raise gr.Error("Invalid token: {e}") from e

    if not api.repo_exists(space_id, repo_type="space"):
        raise gr.Error(f"Space does not exist: 'https://huggingface.co/spaces/{space_id}'.")

    # 1. Add to requirements.txt
    if api.file_exists(repo_id=space_id, repo_type="space", filename="requirements.txt"):
        requirements_file = api.hf_hub_download(repo_id=space_id, repo_type="space", filename="requirements.txt")
        with open(requirements_file) as f:
            requirements = f.read()
    else:
        requirements = ""
    if "gradio-space-ci" not in requirements:
        requirements += "\ngradio-space-ci @ git+https://huggingface.co/spaces/Wauplin/gradio-space-ci@0.2.1\n"

    # 2. Configure CI in README.md
    card = SpaceCard.load(api.hf_hub_download(repo_id=space_id, repo_type="space", filename="README.md"))
    card.data["space_ci"] = {
        "trusted_authors": [],
        "secrets": [],
        "hardware": "cpu-basic",
        "storage": None,
    }
    if card.data.sdk != "gradio":
        raise gr.Error(f"Space must be a Gradio Space, not '{card.data.sdk}'.")
    app_file = card.data.app_file
    if app_file is None:
        raise gr.Error("Space must have an app_file defined their README.")
    if not api.file_exists(repo_id=space_id, repo_type="space", filename=app_file):
        raise gr.Error(f"Could not find app file '{app_file}' in Space repo.")

    # 3. Enable CI in app.py
    with open(api.hf_hub_download(repo_id=space_id, repo_type="space", filename=app_file)) as f:
        app = f.read()
    if "enable_space_ci()" in app:
        raise gr.Error("Space CI is already enabled.")

    all_imports = list(IMPORTS_REGEX.finditer(app))
    if len(all_imports) == 0:
        raise gr.Error("Could not find any imports in app.py.")
    last_import = all_imports[-1]

    app = (
        app[: last_import.end()]
        + "\n\n"
        + "from gradio_space_ci import enable_space_ci"
        + "\n\n"
        + "enable_space_ci()"
        + "\n\n"
        + app[last_import.end() :]
    )

    # 4. Push changes as a PR
    commit = api.create_commit(
        repo_id=space_id,
        repo_type="space",
        operations=[
            CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=str(card).encode()),
            CommitOperationAdd(path_in_repo="requirements.txt", path_or_fileobj=requirements.encode()),
            CommitOperationAdd(path_in_repo=app_file, path_or_fileobj=app.encode()),
        ],
        commit_message="Enable Space CI",
        commit_description=PR_DESCRIPTION,
        create_pr=True,
    )
    assert commit.pr_url is not None  # since `create_pr=True`

    return SUCCESS_MESSAGE.format(space_id=space_id, pr_url=commit.pr_url)