scripts/push_hf_secrets.py

"""
Push all variables from a .env file into a Hugging Face Space as secrets (or variables).

Requirements:
    - huggingface_hub (Python SDK)
        Install: pip install -U huggingface_hub

Usage examples:
    python scripts/push_hf_secrets.py --repo your-username/your-space
    python scripts/push_hf_secrets.py --repo your-username/your-space --env .env.production
    python scripts/push_hf_secrets.py --repo your-username/your-space --dry-run
    python scripts/push_hf_secrets.py --repo your-username/your-space --as-variables  # send as public variables

Notes:
    - This script is intentionally simple and cross-platform.
    - It parses common .env formats (KEY=VALUE, supports quoted values and export prefix).
    - It won’t print secret values; only key names are logged.
    - "Secrets" are private; "Variables" are public. See: Settings → Secrets and variables
"""

from __future__ import annotations

import argparse
import os
import re
import sys
from typing import Dict, Tuple


ENV_LINE_RE = re.compile(r"^\s*(?:export\s+)?([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(.*)\s*$")


def _unquote(value: str) -> str:
    """Strip matching single or double quotes and unescape simple escapes for double quotes.

    - If value is wrapped in double quotes, unescape common sequences (\\n, \\r, \\t, \\" , \\\\).
    - If wrapped in single quotes, return inner content as-is (no escapes processing).
    - Otherwise, return value trimmed of surrounding whitespace.
    """
    if len(value) >= 2 and value[0] == value[-1] and value[0] in ("'", '"'):
        quote = value[0]
        inner = value[1:-1]
        if quote == '"':
            # Process simple escape sequences
            inner = (
                inner.replace(r"\\n", "\n")
                     .replace(r"\\r", "\r")
                     .replace(r"\\t", "\t")
                     .replace(r"\\\"", '"')
                     .replace(r"\\\\", "\\")
            )
        return inner
    return value.strip()


def parse_env_file(path: str) -> Dict[str, str]:
    """Parse a .env-like file into a dict of {KEY: VALUE}.

    Skips blank lines and comments (lines starting with #, ignoring leading whitespace).
    Supports lines like:
      - KEY=VALUE
      - export KEY=VALUE
    Values can be quoted with single or double quotes.
    """
    if not os.path.isfile(path):
        raise FileNotFoundError(f".env file not found: {path}")

    env: Dict[str, str] = {}
    with open(path, "r", encoding="utf-8-sig") as f:
        for idx, raw in enumerate(f, start=1):
            line = raw.rstrip("\n\r")
            stripped = line.strip()
            if not stripped or stripped.startswith("#"):
                continue

            m = ENV_LINE_RE.match(line)
            if not m:
                # Non-fatal: skip lines that don't match KEY=VALUE
                continue

            key, raw_val = m.group(1), m.group(2).strip()

            # If value is unquoted, do not strip inline comments aggressively to avoid breaking tokens.
            value = _unquote(raw_val)
            env[key] = value

    return env


def get_hf_api():
    """Return an authenticated HfApi client or None with a helpful error.

    Uses locally saved token if you previously ran `huggingface-cli login` or
    set HF_TOKEN environment variable.
    """
    try:
        from huggingface_hub import HfApi
    except Exception:
        sys.stderr.write(
            "huggingface_hub is not installed. Install with: pip install -U huggingface_hub\n"
        )
        return None
    return HfApi()

def set_secret(api, repo: str, key: str, value: str, dry_run: bool = False) -> int:
    if dry_run:
        print(f"[DRY RUN] Set secret: {key} -> (hidden) on {repo}")
        return 0
    try:
        api.add_space_secret(repo_id=repo, key=key, value=value)
        print(f"Set secret: {key}")
        return 0
    except Exception as e:
        sys.stderr.write(f"Error setting secret {key!r} for repo {repo!r}: {e}\n")
        return 1


def set_variable(api, repo: str, key: str, value: str, dry_run: bool = False) -> int:
    if dry_run:
        print(f"[DRY RUN] Set variable: {key} -> (hidden) on {repo}")
        return 0
    try:
        api.add_space_variable(repo_id=repo, key=key, value=value)
        print(f"Set variable: {key}")
        return 0
    except Exception as e:
        sys.stderr.write(f"Error setting variable {key!r} for repo {repo!r}: {e}\n")
        return 1


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Push .env variables to a Hugging Face Space as secrets or variables.")
    parser.add_argument("--repo", required=True, help="Space repo id, e.g. your-username/your-space")
    parser.add_argument("--env", default=".env", help="Path to .env file (default: .env)")
    parser.add_argument("--dry-run", action="store_true", help="Print what would be set without applying changes")
    parser.add_argument(
        "--as-variables",
        action="store_true",
        help="Send entries as public variables instead of private secrets",
    )
    parser.add_argument(
        "--exclude",
        action="append",
        default=[],
        help="Key(s) to exclude (can be repeated)",
    )
    args = parser.parse_args(argv)

    api = get_hf_api()
    if api is None:
        return 127

    try:
        env_map = parse_env_file(args.env)
    except Exception as e:
        sys.stderr.write(f"Failed to read env file {args.env}: {e}\n")
        return 2

    if not env_map:
        print("No variables found in .env; nothing to do.")
        return 0

    excluded = set(args.exclude or [])
    total = 0
    failures = 0
    for key, value in env_map.items():
        if key in excluded:
            continue
        total += 1
        if args.as_variables:
            rc = set_variable(api, args.repo, key, value, args.dry_run)
        else:
            rc = set_secret(api, args.repo, key, value, args.dry_run)
        if rc != 0:
            failures += 1

    if failures:
        sys.stderr.write(f"Completed with {failures}/{total} failures.\n")
        return 1

    print(f"Completed: {total} secrets {'validated' if args.dry_run else 'set'} for {args.repo}.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())