vibe/scripts/prepare_release.py

#!/usr/bin/env -S uv run python

from __future__ import annotations

import argparse
from pathlib import Path
import re
import subprocess
import sys
from typing import Any, cast

import tomlkit
from tomlkit.items import Array


def run_git_command(
    *args: str, check: bool = True, capture_output: bool = False
) -> subprocess.CompletedProcess[str]:
    """Run a git command and return the result."""
    result = subprocess.run(
        ["git"] + list(args), check=check, capture_output=capture_output, text=True
    )
    return result


def ensure_public_remote() -> None:
    result = run_git_command("remote", "-v", capture_output=True, check=False)
    remotes = result.stdout

    public_remote_url = "git@github.com:mistralai/mistral-vibe.git"
    if public_remote_url in remotes:
        print("Public remote already exists with correct URL")
        return

    print(f"Creating public remote: {public_remote_url}")
    run_git_command("remote", "add", "public", public_remote_url)
    print("Public remote created successfully")


def switch_to_tag(version: str) -> None:
    tag = f"v{version}"
    print(f"Switching to tag {tag}...")

    result = run_git_command(
        "rev-parse", "--verify", tag, capture_output=True, check=False
    )
    if result.returncode != 0:
        raise ValueError(f"Tag {tag} does not exist")

    run_git_command("switch", "--detach", tag)
    print(f"Successfully switched to tag {tag}")


def get_version_from_pyproject() -> str:
    pyproject_path = Path("pyproject.toml")
    if not pyproject_path.exists():
        raise FileNotFoundError("pyproject.toml not found in current directory")

    content = pyproject_path.read_text()
    version_match = re.search(r'^version = "([^"]+)"$', content, re.MULTILINE)
    if not version_match:
        raise ValueError("Version not found in pyproject.toml")

    return version_match.group(1)


def get_latest_version() -> str:
    result = run_git_command("ls-remote", "--tags", "public", capture_output=True)
    remote_tags_output = (
        result.stdout.strip().split("\n") if result.stdout.strip() else []
    )

    if not remote_tags_output:
        raise ValueError("No version tags found on public remote")

    versions: list[tuple[int, int, int, str]] = []
    MIN_PARTS_IN_LS_REMOTE_LINE = 2  # hash and ref
    for line in remote_tags_output:
        parts = line.split()
        if len(parts) < MIN_PARTS_IN_LS_REMOTE_LINE:
            continue

        _hash, tag_ref = parts[0], parts[1]
        if not tag_ref.startswith("refs/tags/"):
            continue

        tag = tag_ref.replace("refs/tags/", "")
        match = re.match(r"^v(\d+\.\d+\.\d+)$", tag)
        if not match:
            continue

        tag_version = match.group(1)
        try:
            major, minor, patch = parse_version(tag_version)
            versions.append((major, minor, patch, tag_version))
        except ValueError:
            continue

    if not versions:
        raise ValueError(
            "No valid version tags found on public remote (format: vX.Y.Z)"
        )

    versions.sort()

    return max(versions)[3]


def get_base_version_from_current_branch() -> str:
    result = run_git_command(
        "describe",
        "--tags",
        "--match",
        "v[0-9]*.[0-9]*.[0-9]*",
        "--exclude",
        "*-private",
        "--abbrev=0",
        "HEAD",
        capture_output=True,
    )
    tag = result.stdout.strip()
    match = re.match(r"^v(\d+\.\d+\.\d+)$", tag)
    if not match:
        raise ValueError(f"Invalid base version tag found: {tag}")

    return match.group(1)


def parse_version(version_str: str) -> tuple[int, int, int]:
    match = re.match(r"^(\d+)\.(\d+)\.(\d+)$", version_str.strip())
    if not match:
        raise ValueError(f"Invalid version format: {version_str}")

    return int(match.group(1)), int(match.group(2)), int(match.group(3))


def create_release_branch(version: str) -> None:
    branch_name = f"release/v{version}"
    print(f"Creating release branch: {branch_name}")

    result = run_git_command(
        "branch", "--list", branch_name, capture_output=True, check=False
    )
    if result.stdout.strip():
        print(f"Warning: Branch {branch_name} already exists", file=sys.stderr)
        response = input(f"Delete and recreate {branch_name}? (y/N): ")
        if response.lower() == "y":
            run_git_command("branch", "-D", branch_name)
        else:
            print("Aborting", file=sys.stderr)
            sys.exit(1)

    run_git_command("switch", "-c", branch_name)
    print(f"Created and switched to branch {branch_name}")


def cherry_pick_commits(previous_private_tag: str, current_private_tag: str) -> None:
    result = run_git_command(
        "rev-parse", "--verify", previous_private_tag, capture_output=True, check=False
    )
    if result.returncode != 0:
        raise ValueError(f"Tag {previous_private_tag} does not exist")

    result = run_git_command(
        "rev-parse", "--verify", current_private_tag, capture_output=True, check=False
    )
    if result.returncode != 0:
        raise ValueError(f"Tag {current_private_tag} does not exist")

    print(
        f"Cherry-picking commits from {previous_private_tag}..{current_private_tag}..."
    )
    run_git_command("cherry-pick", f"{previous_private_tag}..{current_private_tag}")
    print("Successfully cherry-picked all commits")


def squash_commits(
    previous_version: str,
    current_version: str,
    previous_private_tag: str,
    current_private_tag: str,
) -> None:
    print("Squashing commits into a single release commit...")
    run_git_command("reset", "--soft", f"v{previous_version}")

    # Get all contributors between previous and current private tags
    result = run_git_command(
        "log",
        f"{previous_private_tag}..{current_private_tag}",
        "--format=%aN <%aE>",
        capture_output=True,
    )
    contributors = result.stdout.strip().split("\n")

    # Get current user
    current_user_result = run_git_command("config", "user.email", capture_output=True)
    current_user_email = current_user_result.stdout.strip()

    # Filter out current user and create co-authored lines
    vibe_marker = "vibe@mistral.ai"
    unique_coauthors = {
        f"Co-authored-by: {contributor}"
        for contributor in contributors
        if contributor
        and current_user_email not in contributor
        and vibe_marker not in contributor
    }

    # Add Mistral Vibe as co-author
    coauthored_lines = sorted(unique_coauthors) + [
        "Co-authored-by: Mistral Vibe <vibe@mistral.ai>"
    ]

    # Create commit message
    commit_message = f"v{current_version}\n"
    for line in coauthored_lines:
        commit_message += f"\n{line}"

    # Create the commit
    run_git_command("commit", "-m", commit_message)
    print("Successfully created release commit with co-authors")


def get_pinned_dependencies(group: str | None = None) -> list[str]:
    cmd = [
        "uv",
        "export",
        "--no-hashes",
        "--no-emit-project",
        "--frozen",
        "--format",
        "requirements.txt",
        "--no-annotate",
        "--no-header",
    ]
    if group is None:
        cmd.append("--no-dev")
    else:
        cmd += ["--only-group", group]

    result = subprocess.run(cmd, check=True, capture_output=True, text=True)

    pins: list[str] = []
    for raw_line in result.stdout.splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        pins.append(line)

    if not pins:
        target = group or "[project].dependencies"
        raise ValueError(f"uv export returned no dependencies for {target}")

    return pins


def make_multiline_array(pins: list[str]) -> Array:
    arr = tomlkit.array()
    arr.extend(pins)
    arr.multiline(True)
    return arr


def pin_dependencies(version: str) -> None:
    print("Pinning dependencies for release...")

    pyproject_path = Path("pyproject.toml")
    if not pyproject_path.exists():
        raise FileNotFoundError("pyproject.toml not found in current directory")

    project_pins = get_pinned_dependencies()
    build_pins = get_pinned_dependencies(group="build")

    doc = tomlkit.parse(pyproject_path.read_text())
    cast(dict[str, Any], doc["project"])["dependencies"] = make_multiline_array(
        project_pins
    )
    cast(dict[str, Any], doc["dependency-groups"])["build"] = make_multiline_array(
        build_pins
    )
    pyproject_path.write_text(tomlkit.dumps(doc))
    print(
        f"Pinned {len(project_pins)} project deps and "
        f"{len(build_pins)} build-group deps in pyproject.toml"
    )

    print("Refreshing uv.lock...")
    subprocess.run(["uv", "lock"], check=True)

    run_git_command("add", "pyproject.toml", "uv.lock")
    run_git_command(
        "commit", "--allow-empty", "-m", f"chore: pin dependencies for v{version}"
    )
    print(f"Committed pinned dependencies for v{version}")


def get_commits_summary(previous_version: str, current_version: str) -> str:
    previous_tag = f"v{previous_version}-private"
    current_tag = f"v{current_version}-private"

    result = run_git_command(
        "log", f"{previous_tag}..{current_tag}", "--oneline", capture_output=True
    )
    return result.stdout.strip()


def get_changelog_entry(version: str) -> str:
    changelog_path = Path("CHANGELOG.md")
    if not changelog_path.exists():
        return "CHANGELOG.md not found"

    content = changelog_path.read_text()

    pattern = rf"^## \[{re.escape(version)}\] - .+?(?=^## \[|\Z)"
    match = re.search(pattern, content, re.MULTILINE | re.DOTALL)

    if not match:
        return f"No changelog entry found for version {version}"

    return match.group(0).strip()


def print_summary(
    current_version: str,
    previous_version: str,
    commits_summary: str,
    changelog_entry: str,
    squash: bool,
) -> None:
    print("\n" + "=" * 80)
    print("RELEASE PREPARATION SUMMARY")
    print("=" * 80)
    print(f"\nVersion: {current_version}")
    print(f"Previous version: {previous_version}")
    print(f"Release branch: release/v{current_version}")

    print("\n" + "-" * 80)
    print("COMMITS IN THIS RELEASE")
    print("-" * 80)
    if commits_summary:
        print(commits_summary)
    else:
        print("No commits found")

    print("\n" + "-" * 80)
    print("CHANGELOG ENTRY")
    print("-" * 80)
    print(changelog_entry)

    print("\n" + "-" * 80)
    if not squash:
        print("NEXT STEPS")
        print("-" * 80)
        print(
            f"To review/edit commits before publishing, use interactive rebase:\n"
            f"  git rebase -i v{previous_version}"
        )

        print("\n" + "-" * 80)
    print("REMINDERS")
    print("-" * 80)
    print("Before publishing the release:")
    print("  ✓ Credit any public contributors in the release notes")
    print("  ✓ Close related issues once the release is published")
    print(
        "  ✓ Review and update the changelog if needed "
        "(should be made in the private main branch)"
    )
    print("\n" + "=" * 80)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Prepare a release branch by cherry-picking from private tags",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument("version", help="Version to prepare release for (e.g., 1.1.3)")
    parser.add_argument(
        "--no-squash",
        action="store_false",
        dest="squash",
        default=True,
        help="Disable squashing of commits into a single release commit",
    )
    parser.add_argument(
        "--resume", action="store_true", help="Resume after cherry-picking commits"
    )

    args = parser.parse_args()
    current_version = args.version
    squash = args.squash

    try:
        # Step 1: Ensure public remote exists
        ensure_public_remote()

        if args.resume:
            previous_version = get_base_version_from_current_branch()
        else:
            # Step 2: Fetch all remotes
            print("Fetching all remotes...")
            run_git_command("fetch", "--all")
            print("Successfully fetched all remotes")

            # Step 3: Find latest version
            previous_version = get_latest_version()
        print(f"Previous version: {previous_version}")

        # Step 4: Verify version matches pyproject.toml
        pyproject_version = get_version_from_pyproject()
        if current_version != pyproject_version:
            raise ValueError(
                f"Version mismatch: provided version '{current_version}' does not match "
                f"pyproject.toml version '{pyproject_version}'"
            )
        print(f"Version verified: {current_version}")

        previous_private_tag = f"v{previous_version}-private"
        current_private_tag = f"v{current_version}-private"

        if not args.resume:
            # Step 5: Switch to previous version tag
            switch_to_tag(previous_version)

            # Step 6: Create release branch
            create_release_branch(current_version)

            # Step 7: Cherry-pick commits
            cherry_pick_commits(previous_private_tag, current_private_tag)

        # Step 8: Pin dependencies from uv.lock so the published wheel
        # has frozen Requires-Dist metadata. When squashing, this commit
        # is folded into the release commit; otherwise it stays as the
        # final commit on the release branch.
        pin_dependencies(current_version)

        # Step 9: Squash commits
        if squash:
            squash_commits(
                previous_version,
                current_version,
                previous_private_tag,
                current_private_tag,
            )

        # Step 10: Get summary information
        commits_summary = get_commits_summary(previous_version, current_version)
        changelog_entry = get_changelog_entry(current_version)

        # Step 9: Print summary
        print_summary(
            current_version, previous_version, commits_summary, changelog_entry, squash
        )

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()