JJ triage command (initial version)

jj-triage.py

#!/usr/bin/env python3
import argparse
import textwrap

from dataclasses import dataclass
import subprocess
import re

parser = argparse.ArgumentParser(
    prog="JJ triage",
    usage=textwrap.dedent("""
    Call jj split over and over on remaining changes until we quit.

    When describing selected changes, one can provide triage commands
    at the head of the description like as follows:

    !!abandon

    - abandon the selected changes

    !!squash into target

    - squash the selected changes into target
    """),
    formatter_class=argparse.RawTextHelpFormatter
)

parser.add_argument(
    "-r", "--revision",
    default="@",
    help=textwrap.dedent("""
    The revision from which to start splitting
    (defaults to @).

    On subsequent changes splits will target the remaining
    changes from the previous iteration
    """)
)

# the output of jj split gives us the selected changes and the remaining changes
changes_from_split_regex = re.compile("""Selected changes : (?P<selected_changes>[a-z0-9]+) .+
Remaining changes: (?P<remaining_changes>[a-z0-9]+)""")

def do_split(commit_id) -> tuple[str, str] | None:
    """
    Do a split, returning the selected_changes, remaining_changes
    result pair if a split happened, otherwise return None
    """
    result = subprocess.run(
        "jj split --color never".split(" ") + ["-r", commit_id],
        stderr=subprocess.PIPE
    )
    if result.returncode != 0:
        return None

    stderr_s = result.stderr.decode("utf-8")
    if matched := changes_from_split_regex.search(stderr_s):
        selected_changes, remaining_changes = matched.groups()
        return (selected_changes, remaining_changes)
    else:
        print(repr(stderr_s))
        print(stderr_s)
        raise ValueError("Unexpected split result output!")

def change_description(commit_id: str) -> str:
    """
    Get the description for a change
    """
    result = subprocess.run(["jj", "log", "--no-graph", "-r", commit_id, "-T", "description"], check=True, stdout=subprocess.PIPE)
    return result.stdout.decode("utf-8")

@dataclass
class AbandonAction:
    target_id: str

@dataclass
class SquashAction:
    source_id: str
    target_id: str

# TODO consider how to support absorb nicely
type TriageAction = AbandonAction | SquashAction

def parse_triage_action(commit_id) -> TriageAction | None:
    """
    Parse out triage action if it exists

    !!abandon
    !!squash into abcdef
    """

    description = change_description(commit_id)
    header = description.strip().split("\n")[0]

    if header.startswith("!!"):
        # strip the first bit
        cmd_parts = header[2:].split()
        match cmd_parts[0]:
            case "abandon":
                return AbandonAction(commit_id)
            case "squash":
                if len(cmd_parts) < 3:
                    raise ValueError("Invalid squash command")
                assert cmd_parts[1] == "into", "invalid squash command"
                return SquashAction(
                    source_id=commit_id,
                    target_id=cmd_parts[2]
                )
            case _:
                raise ValueError(f"Unknown command {cmd_parts[0]}")
    return None

def main():
    args = parser.parse_args()
    remaining_changes = args.revision
    # NOTE: do_split returns none if I exit the splitting command
    while split_result := do_split(remaining_changes):
        selected_changes, remaining_changes = split_result
        match parse_triage_action(selected_changes):
            case AbandonAction(commit_id):
                subprocess.run(["jj", "abandon", commit_id], check=True)
            case SquashAction(source_id, target_id):
                # NOTE -u to drop the source message (which is just a !!squash notation)
                subprocess.run(["jj", "squash", "--from", source_id, "--into", target_id, "-u"], check=True)
            case None:
                pass

if __name__ == "__main__":
   main()