codeflash-agent/scripts/versioning.py

"""Version management for workspace packages.

Usage:
    python scripts/versioning.py check-version          # auto-detect changed packages
    python scripts/versioning.py check-version PKG       # check one package
    python scripts/versioning.py version-dev PKG         # bump to pre-release
    python scripts/versioning.py version-release PKG     # release current version
    python scripts/versioning.py get-version PKG         # print current version
"""

from __future__ import annotations

import subprocess
import sys
from pathlib import Path

import tomlkit

REPO_ROOT = Path(__file__).resolve().parent.parent
PACKAGES_DIR = REPO_ROOT / "packages"

# Packages that participate in versioning (have meaningful releases).
VERSIONED_PACKAGES = ["codeflash-core", "codeflash-python"]


def _read_version(pyproject_path: Path) -> str:
    """Read version string from a pyproject.toml."""
    data = tomlkit.parse(pyproject_path.read_text())
    return str(data["project"]["version"])


def _write_version(pyproject_path: Path, new_version: str) -> None:
    """Write a new version string into a pyproject.toml."""
    data = tomlkit.parse(pyproject_path.read_text())
    data["project"]["version"] = new_version
    pyproject_path.write_text(tomlkit.dumps(data))


def _bump_patch(version: str) -> str:
    """Bump the patch component: 0.1.0 -> 0.1.1, 0.1.1.dev0 -> 0.1.2."""
    base = version.split(".dev", maxsplit=1)[0]
    parts = base.split(".")
    parts[-1] = str(int(parts[-1]) + 1)
    return ".".join(parts)


def _current_branch() -> str:
    """Return the current git branch name."""
    result = subprocess.run(
        ["git", "rev-parse", "--abbrev-ref", "HEAD"],
        capture_output=True,
        text=True,
        check=True,
    )
    return result.stdout.strip()


def _changed_packages() -> list[str]:
    """Return package names that have file changes vs origin/main."""
    changed = []
    for pkg in VERSIONED_PACKAGES:
        pkg_path = f"packages/{pkg}/"
        result = subprocess.run(
            [
                "git",
                "diff",
                "--name-only",
                "origin/main..HEAD",
                "--",
                pkg_path,
            ],
            capture_output=True,
            text=True,
        )
        if result.stdout.strip():
            changed.append(pkg)
    return changed


# -- Commands ----------------------------------------------------------------


def cmd_get_version(package: str) -> None:
    """Print the current version of a package."""
    pyproject = PACKAGES_DIR / package / "pyproject.toml"
    if not pyproject.exists():
        print(f"No pyproject.toml found at {pyproject}")
        sys.exit(1)
    print(_read_version(pyproject))


def cmd_check_version(package: str | None = None) -> None:
    """Verify that changed packages have bumped their version vs origin/main.

    If *package* is None, auto-detect which packages changed.
    """
    packages = [package] if package else _changed_packages()

    if not packages:
        print("No versioned packages changed.")
        return

    failed = False
    for pkg in packages:
        pyproject_rel = f"packages/{pkg}/pyproject.toml"
        pyproject = PACKAGES_DIR / pkg / "pyproject.toml"

        if not pyproject.exists():
            print(f"No pyproject.toml at {pyproject}")
            failed = True
            continue

        current = _read_version(pyproject)

        # Get the version on origin/main.
        result = subprocess.run(
            ["git", "show", f"origin/main:{pyproject_rel}"],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            print(f"{pkg}: new package, no version on main yet. OK.")
            continue

        main_data = tomlkit.parse(result.stdout)
        main_version = str(main_data["project"]["version"])

        if current == main_version:
            print(
                f"{pkg}: version {current} unchanged from main. "
                "Run `make version-dev ARGS={pkg}` to bump."
            )
            failed = True
        else:
            print(f"{pkg}: {main_version} -> {current}. OK.")

    if failed:
        sys.exit(1)


def cmd_version_dev(package: str) -> None:
    """Bump to pre-release version and create a changelog entry."""
    branch = _current_branch()
    if branch == "main":
        print("Cannot bump version on main branch.")
        sys.exit(1)

    pyproject = PACKAGES_DIR / package / "pyproject.toml"
    if not pyproject.exists():
        print(f"No pyproject.toml at {pyproject}")
        sys.exit(1)

    current = _read_version(pyproject)
    new_version = _bump_patch(current) + ".dev0"
    _write_version(pyproject, new_version)
    print(f"{package}: {current} -> {new_version}")

    # Create changelog entry.
    changelogs_dir = PACKAGES_DIR / package / "changelogs"
    changelogs_dir.mkdir(exist_ok=True)
    safe_branch = branch.replace("_", "-").replace("/", "-")
    entry_path = changelogs_dir / f"{safe_branch}.md"
    entry_path.write_text("### Enhancements\n\n### Features\n\n### Fixes\n")
    subprocess.run(["git", "add", str(entry_path)], check=True)
    print(f"Created changelog entry: {entry_path.relative_to(REPO_ROOT)}")


def cmd_version_release(package: str) -> None:
    """Release the current version (strip .dev suffix or bump patch)."""
    pyproject = PACKAGES_DIR / package / "pyproject.toml"
    if not pyproject.exists():
        print(f"No pyproject.toml at {pyproject}")
        sys.exit(1)

    current = _read_version(pyproject)
    if ".dev" in current:
        new_version = current.split(".dev")[0]
    else:
        new_version = _bump_patch(current)

    _write_version(pyproject, new_version)
    print(f"{package}: released {new_version}")


# -- CLI dispatch ------------------------------------------------------------


def main() -> None:
    """Dispatch subcommand."""
    usage = (
        "Usage: python scripts/versioning.py "
        "<check-version|version-dev|version-release|get-version> [PACKAGE]"
    )

    if len(sys.argv) < 2:
        print(usage)
        sys.exit(1)

    command = sys.argv[1]
    package = sys.argv[2] if len(sys.argv) > 2 else None

    if command == "check-version":
        cmd_check_version(package)
    elif command == "get-version":
        if not package:
            print("get-version requires a package name.")
            sys.exit(1)
        cmd_get_version(package)
    elif command == "version-dev":
        if not package:
            print("version-dev requires a package name.")
            sys.exit(1)
        cmd_version_dev(package)
    elif command == "version-release":
        if not package:
            print("version-release requires a package name.")
            sys.exit(1)
        cmd_version_release(package)
    else:
        print(f"Unknown command: {command}")
        print(usage)
        sys.exit(1)


if __name__ == "__main__":
    main()