git-rebase-script

#!/usr/bin/env python3

import os
import shlex
import subprocess
import sys
from pathlib import Path


SCRIPT_PATH = os.path.abspath(sys.argv[0])


def run_git(*args, capture=True, check=True):
    """Run a git command and return stdout."""
    cmd = ["git"] + list(args)
    if capture:
        result = subprocess.run(cmd, capture_output=True, text=True, check=check)
        return result.stdout.strip()
    else:
        subprocess.run(cmd, check=check)


def git_rev_parse(ref):
    """Get the full hash for a reference."""
    return run_git("rev-parse", ref)


def git_rev_parse_short(ref):
    """Get the short hash for a reference."""
    return run_git("rev-parse", "--short", ref)


def resolve_ref(ref):
    """Resolve a reference to a commit hash.
    
    First tries to resolve as a git reference (hash, branch, HEAD~n).
    If that fails, searches for a commit with matching message in current HEAD.
    """
    # Try as a git reference first
    try:
        return git_rev_parse(ref)
    except subprocess.CalledProcessError:
        pass
    
    # Search for commit by message (only in current HEAD)
    try:
        matches = run_git("log", "--format=%H", "--fixed-strings", f"--grep={ref}", "HEAD")
        if not matches:
            print(f"Error: no commit found matching '{ref}'", file=sys.stderr)
            sys.exit(1)
        
        lines = matches.split('\n')
        if len(lines) > 1:
            print(f"Error: multiple commits match '{ref}':", file=sys.stderr)
            for line in lines[:5]:  # Show first 5 matches
                msg = run_git("log", "-1", "--format=%s", line)
                print(f"  {git_rev_parse_short(line)} {msg}", file=sys.stderr)
            if len(lines) > 5:
                print(f"  ... and {len(lines) - 5} more", file=sys.stderr)
            sys.exit(1)
        
        return lines[0]
    except subprocess.CalledProcessError:
        print(f"Error: could not resolve reference '{ref}'", file=sys.stderr)
        sys.exit(1)


def is_ancestor(older, newer):
    """Return True if older is an ancestor of newer."""
    try:
        run_git("merge-base", "--is-ancestor", older, newer)
        return True
    except subprocess.CalledProcessError:
        return False


def usage():
    """Print usage and exit."""
    prog = Path(sys.argv[0]).name
    print(f"Usage:")
    print(f"  {prog} edit HASH|MESSAGE")
    print(f"  {prog} reword MESSAGE HASH|MESSAGE")
    print(f"  {prog} squash MESSAGE HASH|MESSAGE...")
    print(f"  {prog} move REF_HASH|MESSAGE HASH|MESSAGE...")
    print()
    print(f"Examples:")
    print(f"  {prog} edit HEAD~3")
    print(f"  {prog} edit 'Fix typo in README'")
    print(f"  {prog} reword 'Fix typo in README' abc1234")
    print(f"  {prog} squash 'Combined feature' 'Add feature X' 'Fix feature X'")
    print(f"  {prog} move 'Initial commit' 'Move this commit'")
    sys.exit(1)


def run_rebase(args, env=None):
    """Run git rebase, exiting with git's status to avoid stack traces."""
    result = subprocess.run(["git", "rebase", "-i", *args], env=env)
    if result.returncode != 0:
        sys.exit(result.returncode)


def editor_reword(filepath, message):
    """Editor function for reword: replace commit message."""
    with open(filepath, 'r') as f:
        content = f.read()

    if content:
        with open(filepath, 'w') as f:
            f.write(message + '\n')


def editor_squash(filepath, message):
    """Editor function for squash: replace combined commit message."""
    with open(filepath, 'r') as f:
        lines = f.readlines()

    if lines and lines[0].startswith('# This is a combination of'):
        with open(filepath, 'w') as f:
            f.write(message + '\n')
            for line in lines:
                if line.startswith('#'):
                    f.write(line)


def replace_command(line, command):
    """Replace the leading rebase command on a todo line."""
    parts = line.split(maxsplit=1)
    suffix = f" {parts[1]}" if len(parts) > 1 else ""
    newline = '\n' if line.endswith('\n') else ''
    return f"{command}{suffix}{newline}"


def editor_squash_order(filepath, ordered_shorts):
    """Editor for squash: reorder commits so they are consecutive."""
    with open(filepath, 'r') as f:
        lines = f.readlines()

    index_by_short = {}
    line_by_short = {}

    for idx, line in enumerate(lines):
        if not line.startswith('pick '):
            continue
        parts = line.split()
        if len(parts) < 2:
            continue
        commit_short = parts[1]
        if commit_short in ordered_shorts:
            index_by_short[commit_short] = idx
            line_by_short[commit_short] = line

    for short in ordered_shorts:
        if short not in index_by_short:
            print(f"Error: commit {short} not found in rebase todo", file=sys.stderr)
            sys.exit(1)

    first_index = min(index_by_short.values())
    selected_indices = set(index_by_short.values())

    block = []
    block.append(replace_command(line_by_short[ordered_shorts[0]], "pick"))
    for short in ordered_shorts[1:]:
        block.append(replace_command(line_by_short[short], "squash"))

    result = []
    for idx, line in enumerate(lines):
        if idx == first_index:
            result.extend(block)
        if idx in selected_indices:
            continue
        result.append(line)

    with open(filepath, 'w') as f:
        f.writelines(result)


def editor_move(filepath, ref_short, move_shorts):
    """Editor function for move: reorder commits in todo."""
    with open(filepath, 'r') as f:
        lines = f.readlines()

    ref_line = None
    ref_index = None
    move_lines = {}
    other_lines = []

    for line in lines:
        if line.startswith('pick '):
            parts = line.split()
            hash_match = parts[1] if len(parts) > 1 else None
            if hash_match == ref_short:
                ref_line = line
                ref_index = len(other_lines)
                other_lines.append(line)
            elif hash_match in move_shorts:
                move_lines[hash_match] = line
            else:
                other_lines.append(line)
        else:
            other_lines.append(line)

    if ref_line is None:
        print(f"Error: reference commit {ref_short} not found in rebase todo", file=sys.stderr)
        sys.exit(1)

    for short in move_shorts:
        if short not in move_lines:
            print(f"Error: commit {short} not found in rebase todo", file=sys.stderr)
            sys.exit(1)

    result = other_lines[:ref_index + 1]
    for short in move_shorts:
        result.append(move_lines[short])
    result.extend(other_lines[ref_index + 1:])

    with open(filepath, 'w') as f:
        f.writelines(result)


def action_edit(commit):
    """Mark a commit for editing."""
    full = resolve_ref(commit)
    short = git_rev_parse_short(full)

    editor = f"sed -i -e 's/^pick {short}/edit {short}/'"
    env = os.environ.copy()
    env["GIT_SEQUENCE_EDITOR"] = editor

    run_rebase([f"{full}~1"], env)


def action_reword(message, commit):
    """Reword a commit message."""
    full = resolve_ref(commit)
    short = git_rev_parse_short(full)

    sequence_editor = f"sed -i -e 's/^pick {short}/reword {short}/'"
    git_editor = f"{shlex.quote(SCRIPT_PATH)} --editor-reword {shlex.quote(message)}"

    env = os.environ.copy()
    env["GIT_SEQUENCE_EDITOR"] = sequence_editor
    env["GIT_EDITOR"] = git_editor

    run_rebase([f"{full}~1"], env)


def action_squash(message, *commits):
    """Squash commits together (any order)."""
    if not commits:
        usage()

    full_commits = [resolve_ref(c) for c in commits]

    if len(full_commits) != len(set(full_commits)):
        print("Error: duplicate commits provided for squash", file=sys.stderr)
        sys.exit(1)

    oldest_full = None
    newest_full = None

    for candidate in full_commits:
        if all(candidate == other or is_ancestor(candidate, other) for other in full_commits):
            oldest_full = candidate
            break

    for candidate in full_commits:
        if all(candidate == other or is_ancestor(other, candidate) for other in full_commits):
            newest_full = candidate
            break

    if oldest_full is None or newest_full is None:
        print("Error: commits are not on the same linear history", file=sys.stderr)
        sys.exit(1)

    oldest_parent = None
    try:
        oldest_parent = git_rev_parse(f"{oldest_full}~1")
    except subprocess.CalledProcessError:
        oldest_parent = None

    history_range = f"{oldest_parent}..{newest_full}" if oldest_parent else f"{oldest_full}..{newest_full}"
    history = run_git("rev-list", "--reverse", "--ancestry-path", history_range)
    ordered = [c for c in history.splitlines() if c in full_commits]

    if oldest_parent is None:
        ordered = [oldest_full] + ordered

    if len(ordered) != len(full_commits):
        print("Error: commits are not on the same linear history", file=sys.stderr)
        sys.exit(1)

    full_commits = ordered
    short_commits = [git_rev_parse_short(f) for f in full_commits]

    git_editor = f"{shlex.quote(SCRIPT_PATH)} --editor-squash {shlex.quote(message)}"
    editor_args = [shlex.quote(SCRIPT_PATH), '--editor-squash-order']
    editor_args.extend(shlex.quote(s) for s in short_commits)
    sequence_editor = ' '.join(editor_args)

    env = os.environ.copy()
    env["GIT_SEQUENCE_EDITOR"] = sequence_editor
    env["GIT_EDITOR"] = git_editor

    if oldest_parent:
        run_rebase([oldest_parent], env)
    else:
        run_rebase(["--root"], env)


def action_move(ref_hash, *move_hashes):
    """Move commits to appear after a reference commit."""
    if not move_hashes:
        usage()

    if len(move_hashes) != len(set(move_hashes)):
        print(f"Error: duplicate hashes in move list", file=sys.stderr)
        sys.exit(1)

    ref_full = resolve_ref(ref_hash)
    ref_short = git_rev_parse_short(ref_full)

    move_full = [resolve_ref(h) for h in move_hashes]
    move_short = [git_rev_parse_short(f) for f in move_full]

    if len(move_short) != len(set(move_short)):
        print(f"Error: duplicate commits after resolving hashes", file=sys.stderr)
        sys.exit(1)

    rebase_args = None
    try:
        base = git_rev_parse(f"{ref_full}~1")
        rebase_args = [base]
    except subprocess.CalledProcessError:
        # Reference commit is the root; rebase from the start of history.
        rebase_args = ["--root"]

    editor_args = [shlex.quote(SCRIPT_PATH), '--editor-move', shlex.quote(ref_short)]
    editor_args.extend(shlex.quote(s) for s in move_short)
    sequence_editor = ' '.join(editor_args)

    env = os.environ.copy()
    env["GIT_SEQUENCE_EDITOR"] = sequence_editor

    run_rebase(rebase_args, env)


def main():
    if len(sys.argv) < 2:
        usage()

    if sys.argv[1] == "--editor-reword":
        if len(sys.argv) != 4:
            sys.exit(1)
        editor_reword(sys.argv[3], sys.argv[2])
        return

    if sys.argv[1] == "--editor-squash":
        if len(sys.argv) != 4:
            sys.exit(1)
        editor_squash(sys.argv[3], sys.argv[2])
        return

    if sys.argv[1] == "--editor-squash-order":
        if len(sys.argv) < 4:
            sys.exit(1)
        ordered_shorts = sys.argv[2:-1]
        filepath = sys.argv[-1]
        editor_squash_order(filepath, ordered_shorts)
        return

    if sys.argv[1] == "--editor-move":
        if len(sys.argv) < 4:
            sys.exit(1)
        ref_short = sys.argv[2]
        move_shorts = sys.argv[3:-1]
        filepath = sys.argv[-1]
        editor_move(filepath, ref_short, move_shorts)
        return

    if len(sys.argv) < 3:
        usage()

    action = sys.argv[1]
    args = sys.argv[2:]

    if action == "edit":
        if len(args) != 1:
            usage()
        action_edit(args[0])

    elif action == "reword":
        if len(args) != 2:
            usage()
        action_reword(args[0], args[1])

    elif action == "squash":
        if len(args) < 2:
            usage()
        action_squash(args[0], *args[1:])

    elif action == "move":
        if len(args) < 2:
            usage()
        action_move(args[0], *args[1:])

    else:
        usage()


if __name__ == "__main__":
    main()