Manage a stack of commit-groups

git-ladder

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Author: Maarten Nieber
# Url: https://gist.github.com/mnieber/

import json
import subprocess
import sys
from argparse import ArgumentParser
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional


@dataclass
class LadderConfig:
    name: str
    ladder_branch: str
    pr_branch: Optional[str]
    target_branch: str


def run_git(*args, check=True) -> str:
    """Run a git command and return its output."""
    result = subprocess.run(["git"] + list(args), capture_output=True, text=True)
    if check and result.returncode != 0:
        print(f"Git command failed: {result.stderr}", file=sys.stderr)
        sys.exit(1)
    return result.stdout.strip()


def _get_git_root() -> Path:
    """Get the root directory of the git repository."""
    return Path(run_git("rev-parse", "--show-toplevel"))


def _get_current_branch() -> str:
    """Get the name of the current git branch."""
    return run_git("rev-parse", "--abbrev-ref", "HEAD")


def _load_ladder_config(ladder_name: str) -> LadderConfig:
    """Load the configuration for a specific ladder."""
    config_path = (
        _get_git_root()
        / ".git"
        / "ladder"
        / "ladders"
        / ladder_name
        / "ladder-config.json"
    )
    if not config_path.exists():
        print(f"Ladder {ladder_name} not found", file=sys.stderr)
        sys.exit(1)

    with open(config_path) as f:
        data = json.load(f)
        return LadderConfig(
            name=data["name"],
            ladder_branch=data["ladder_branch"],
            pr_branch=data.get("pr_branch"),  # Optional
            target_branch=data["target_branch"],
        )


def _save_ladder_config(config: LadderConfig):
    """Save the ladder configuration to disk."""
    config_path = (
        _get_git_root()
        / ".git"
        / "ladder"
        / "ladders"
        / config.name
        / "ladder-config.json"
    )
    config_path.parent.mkdir(parents=True, exist_ok=True)

    with open(config_path, "w") as f:
        json.dump(
            {
                "name": config.name,
                "ladder_branch": config.ladder_branch,
                "pr_branch": config.pr_branch,
                "target_branch": config.target_branch,
            },
            f,
        )


def _create_marker_commit(rung_name: str, ladder_name: str, begin: bool):
    """Create a marker commit for a rung (start or end)."""
    marker_path = (
        _get_git_root() / ".git-ladder" / "ladders" / ladder_name / "rungs" / rung_name
    )

    if begin:
        marker_path.parent.mkdir(parents=True, exist_ok=True)
        marker_path.touch()
        run_git("add", str(marker_path))
        run_git("commit", "-m", _get_rung_commit_name(rung_name, begin=True))
    else:
        run_git("rm", str(marker_path))
        run_git("commit", "-m", _get_rung_commit_name(rung_name, begin=False))


def _get_rung_commit_name(rung_name, begin):
    prefix = "B:" if begin else "E:"
    label = f" {prefix} {rung_name}"
    nr_of_dashes = 99 - len(label)
    return "-" * nr_of_dashes + label


def _get_rung_commit_sha(first_rung, begin):
    return run_git(
        "log",
        "--grep",
        _get_rung_commit_name(first_rung, begin=begin),
        "--format=%H",
    )


def _get_active_ladder(raise_on_error=True):
    branch_name = _get_current_branch()
    if branch_name.startswith("ladder/"):
        return branch_name[len("ladder/") :]

    if raise_on_error:
        print("No active ladder", file=sys.stderr)
        sys.exit(1)

    return ""


def _get_rungs() -> List[str]:
    """
    Extract rung names from git commit history by looking for special marker commits.
    Returns a list of rung names in chronological order.
    """
    # Get all commits that match our marker pattern (both begin and end markers)
    git_log = run_git(
        "log",
        "--grep=^-\\{20,\\} [BE]:",  # Match 20+ dashes followed by B: or E:
        "--format=%s",  # Only get commit subjects
        "--reverse",  # Get oldest commits first
    )

    if not git_log:
        return []

    rungs = []
    for commit in git_log.split("\n"):
        # Extract the rung name from commit message
        # Format is "----...----[B|E]: rung_name"
        marker_type = (
            commit[-len(commit) : commit.rfind(":")].strip("-").strip(" ")
        )  # Get 'B' or 'E'
        rung_name = commit[commit.rfind(":") + 1 :].strip()

        if marker_type == "B":
            # Only add rung name on begin markers to maintain order
            rungs.append(rung_name)

    return rungs


def cmd_checkout_new_ladder(args):
    ladder_branch = f"ladder/{args.ladder}"

    # check that ladder_branch does not exist
    if run_git("branch", "--list", ladder_branch):
        print(f"Branch {ladder_branch} already exists", file=sys.stderr)
        sys.exit(1)

    current_branch = _get_current_branch()

    # Create and checkout the ladder branch
    run_git("checkout", "-b", ladder_branch)

    # Create initial ladder config
    ladder_config = LadderConfig(
        name=args.ladder,
        ladder_branch=ladder_branch,
        pr_branch=None,
        target_branch=current_branch,
    )
    _save_ladder_config(ladder_config)


def cmd_checkout_ladder(args):
    config = _load_ladder_config(args.existing_ladder)
    run_git("checkout", config.ladder_branch)


def cmd_add_rung(args):
    rung_name = args.name

    # Get active ladder config
    active_ladder = _get_active_ladder()
    ladder_config = _load_ladder_config(active_ladder)

    # Add end marker to previous rung
    rungs = _get_rungs()
    if rungs:
        previous_rung = rungs[-1]
        _create_marker_commit(previous_rung, active_ladder, begin=False)

    # Create start marker for the new rung
    _create_marker_commit(rung_name, active_ladder, begin=True)

    if args.close:
        # Create end marker for the new rung
        _create_marker_commit(rung_name, active_ladder, begin=False)


def cmd_create_pr(args):
    # Get active ladder config
    active_ladder = _get_active_ladder()
    ladder_config = _load_ladder_config(active_ladder)

    rungs = _get_rungs()
    if not rungs:
        print("Ladder has no rungs", file=sys.stderr)
        sys.exit(1)

    if args.branch:
        # Create new PR branch
        if run_git("branch", "--list", args.branch, check=False):
            print(f"Branch {args.branch} already exists", file=sys.stderr)
            sys.exit(1)

        # Find end of first rung
        first_rung = rungs[0]
        end_commit = _get_rung_commit_sha(first_rung, begin=False)

        run_git("branch", args.branch, end_commit)
        ladder_config.pr_branch = args.branch
        _save_ladder_config(ladder_config)
    else:
        # Reset existing PR branch
        if not ladder_config.pr_branch:
            print("No PR branch set for this ladder", file=sys.stderr)
            sys.exit(1)

        first_rung = rungs[0]
        end_commit = _get_rung_commit_sha(first_rung, begin=False)

        if not args.force:
            current = run_git("rev-parse", ladder_config.pr_branch)
            print(
                f"Would reset {ladder_config.pr_branch} from {current} to {end_commit}"
            )
        else:
            run_git("branch", "-f", ladder_config.pr_branch, end_commit)


def cmd_status(args):
    active_ladder = _get_active_ladder()
    ladder_config = _load_ladder_config(active_ladder)

    if ladder_config.pr_branch:
        # Check if PR branch exists and points to end of first rung
        first_rung = _get_rungs()[0]
        end_commit = _get_rung_commit_sha(first_rung, begin=False)

        try:
            pr_commit = run_git("rev-parse", ladder_config.pr_branch, check=False)
            if pr_commit != end_commit:
                print("pr-branch-not-found")
                return
        except:
            print("pr-branch-not-found")
            return


def cmd_restack():
    # Get active ladder config
    active_ladder = _get_active_ladder()
    ladder_config = _load_ladder_config(active_ladder)

    if not ladder_config.pr_branch:
        print("No PR branch set for this ladder", file=sys.stderr)
        sys.exit(1)

    # Get the end of first rung
    first_rung = _get_rungs()[0]
    first_rung_end = _get_rung_commit_sha(first_rung, begin=False)

    # Rebase all commits after the first rung onto the PR branch
    run_git(
        "rebase",
        "--onto",
        ladder_config.pr_branch,
        first_rung_end,
        ladder_config.ladder_branch,
    )


def _parse_args():
    parser = ArgumentParser(description="Manage git ladders")
    subparsers = parser.add_subparsers(dest="command", required=True)

    checkout_parser = subparsers.add_parser("checkout", help="Checkout a ladder")
    checkout_parser.add_argument(
        "existing_ladder",
        help="The name of the existing ladder to checkout",
        nargs="?",
        default=None,
    )
    checkout_parser.add_argument(
        "--ladder", help="The name of the new ladder to create"
    )

    add_rung_parser = subparsers.add_parser("add-rung", help="Add a new rung")
    add_rung_parser.add_argument("name", help="The name of the rung to add")
    add_rung_parser.add_argument(
        "--close", help="Create an end-marker for the new rung"
    )

    create_pr_parser = subparsers.add_parser("create-pr", help="Create a PR branch")
    create_pr_parser.add_argument("branch", nargs="?", help="The name of the PR branch")
    create_pr_parser.add_argument(
        "--force", action="store_true", help="Overwrite existing PR branch"
    )

    subparsers.add_parser("status", help="Show the status of the active ladder")

    subparsers.add_parser("restack", help="Restack the ladder onto the PR branch")

    return parser.parse_args()


def main():
    args = _parse_args()

    commands = {
        "checkout": (
            cmd_checkout_new_ladder
            if (args.command == "checkout" and args.ladder)
            else cmd_checkout_ladder
        ),
        "add-rung": cmd_add_rung,
        "create-pr": cmd_create_pr,
        "status": cmd_status,
        "restack": cmd_restack,
    }

    command = commands[args.command]
    command(args)


if __name__ == "__main__":
    main()