doink.py

#!/usr/bin/env python3
"""
Doink: Development utilities for Sentry projects.

A collection of git-aware development tools that help streamline common workflows:
- peep: Find and run tests for modified files
- modified: List and open modified files  
- otherdir: Navigate between src/sentry and tests/sentry directories
- meep: Run mypy on modified Python files and their tests
- prune: Clean up local branches from closed/merged PRs
- install/uninstall: Manage shell function installation in ~/.zshrc

All commands work with git to automatically detect modified files since the merge base
with the upstream branch, making it easy to focus on your current changes.
"""
import argparse
import json
import os
import re
import shlex
import subprocess
import sys
from pathlib import Path
from typing import Callable, NamedTuple, Sequence


def get_git_root() -> Path:
    return Path(
        subprocess.check_output(["git", "rev-parse", "--show-toplevel"])
        .decode("utf-8")
        .strip()
    )


def get_upstream_branch() -> str:
    """Get the default upstream branch from git."""
    return (
        subprocess.check_output(["git", "symbolic-ref", "refs/remotes/origin/HEAD"])
        .decode("utf-8")
        .strip()
        .removeprefix("refs/remotes/")
    )


def _git_diff_files(args: list[str]) -> list[Path]:
    return [
        Path(f)
        for f in subprocess.check_output(
            ["git", "diff", "--name-only", "--diff-filter=ACM"] + args
        )
        .decode("utf-8")
        .splitlines()
    ]


def get_modified_files(
    cwd: Path, include_committed: bool = True, suffixes: set[str] | None = None
) -> set[Path]:
    upstream_branch = get_upstream_branch()
    merge_base = (
        subprocess.check_output(["git", "merge-base", "HEAD", upstream_branch])
        .decode("utf-8")
        .strip()
    )
    unstaged_files = _git_diff_files([])
    staged_files = _git_diff_files(["--cached"])
    from_upstream = (
        _git_diff_files([f"{merge_base}..HEAD"]) if include_committed else []
    )
    files = set(unstaged_files) | set(staged_files) | set(from_upstream)
    if suffixes:
        files = {f for f in files if f.suffix in suffixes}
    return files


def get_test_file(git_root: Path, file: Path) -> Path | None:
    """
    Returns a git-root-relative path to the test file for the given file.
    """
    if file.suffix != ".py":
        return None
    file = git_root / file
    if not file.exists():
        return None
    file = file.absolute().resolve()
    sentry_root = git_root / "src" / "sentry"
    test_root = git_root / "tests" / "sentry"
    if (
        file.is_relative_to(test_root)
        and file.name.startswith("test_")
        and file.exists()
    ):
        return file
    if file.is_relative_to(sentry_root):
        rel = file.relative_to(git_root)
        test_file = test_root / Path(*rel.with_name("test_" + rel.name).parts[2:])
        if test_file.exists():
            return test_file
    return None


def get_other_path(git_root: Path, curr_dir: Path) -> Path | None:
    """
    If curr_dir is in tests/sentry, returns the corresponding dir in src/sentry.
    If curr_dir is in src/sentry, returns the corresponding dir in tests/sentry.
    If neither, returns None.
    """
    curr_dir = curr_dir.absolute().resolve()
    sentry_root = git_root / "src" / "sentry"
    test_root = git_root / "tests" / "sentry"

    if curr_dir.is_relative_to(sentry_root):
        rel = curr_dir.relative_to(sentry_root)
        return test_root / rel
    elif curr_dir.is_relative_to(test_root):
        rel = curr_dir.relative_to(test_root)
        return sentry_root / rel
    return None


def print_cmd[T](cmd_elts: Sequence[T]) -> None:
    print(shlex.join([str(elt) for elt in cmd_elts]))


def run_command(wd: Path, cmd: Sequence[str | Path], args: argparse.Namespace) -> None:
    if args.dry_run:
        print_cmd(cmd)
    else:
        subprocess.run(cmd, cwd=wd)


def dedup[T](items: list[T]) -> list[T]:
    return list(dict.fromkeys(items))


def peep_command(args: argparse.Namespace) -> None:
    git_root = get_git_root()
    files = args.files or get_modified_files(git_root, include_committed=True)
    test_files = []
    for file in files:
        test_file = get_test_file(git_root, file)
        if test_file:
            test_files.append(test_file)
    if not test_files:
        print("No test files found")
        return
    test_files = dedup(test_files)
    print("Found", len(test_files), "test files for", len(files), "targeted files.")
    if args.open:
        run_command(git_root, ["cursor", "-r", *test_files], args)
    else:
        run_command(git_root, ["pytest", *test_files], args)


def modified_files_command(args: argparse.Namespace) -> None:
    git_root = get_git_root()
    files = get_modified_files(git_root, include_committed=True)
    if args.type:
        suffix = "." + args.type
        files = {f for f in files if f.suffix == suffix}
    if args.open:
        run_command(git_root, ["cursor", "-r", *sorted(files)], args)
    else:
        print(shlex.join(sorted([f.as_posix() for f in files])))


def otherdir_command(args: argparse.Namespace) -> None:
    git_root = get_git_root()
    cwd = Path.cwd()
    other_path = get_other_path(git_root, cwd)
    if other_path and other_path.exists():
        print(os.path.relpath(other_path, cwd))
    else:
        sys.exit("No other path found. Are you in a test or src directory?")


def meep_command(args: argparse.Namespace) -> None:
    git_root = get_git_root()
    files = get_modified_files(git_root, include_committed=True, suffixes={".py"})
    if not files:
        print("No modified Python files found")
        return

    all_files = [*files]
    for file in files:
        test_file = get_test_file(git_root, file)
        if test_file:
            all_files.append(test_file.relative_to(git_root))

    # Convert to repo-relative paths
    rel_paths = [f.as_posix() for f in dedup(all_files)]
    print(f"Running mypy on {len(rel_paths)} files...")

    # Run mypy from repo root without changing caller's cwd
    cmd = ["mypy", *rel_paths]
    run_command(git_root, cmd, args)


def get_short_branch_names() -> list[str]:
    return [
        line.decode("utf-8").strip()
        for line in subprocess.check_output(
            ["git", "branch", "--format=%(refname:short)", "--sort=committerdate"]
        ).splitlines()
    ]


def get_branch_details(branches: list[str]) -> dict[str, tuple[str, str]]:
    """Get the commit message and relative age for multiple branches in a single git call."""
    # Create a format that includes the branch name and relative age
    format_str = "%(refname:short)%00%(creatordate:relative)"
    output = (
        subprocess.check_output(
            [
                "git",
                "for-each-ref",
                "--format=" + format_str,
                *[f"refs/heads/{b}" for b in branches],
            ]
        )
        .decode("utf-8")
        .strip()
        .splitlines()
    )

    result = {}
    for line in output:
        branch, date = line.split("\0")
        result[branch] = date
    return result


def prune_command(args: argparse.Namespace) -> None:
    git_root = get_git_root()
    closed_prs_cmd = [
        "gh",
        "pr",
        "list",
        "--state",
        "closed",
        "--author",
        "@me",
        "--base",
        "master",
        "--json",
        "headRefName,title",
        "-q",
        ".[] | {branch: .headRefName, title: .title}",
    ]
    closed_prs = subprocess.check_output(closed_prs_cmd)
    closed_pr_data = [json.loads(line) for line in closed_prs.splitlines()]
    closed_pr_branches = {pr["branch"] for pr in closed_pr_data}
    pr_titles = {pr["branch"]: pr["title"] for pr in closed_pr_data}

    all_branches = get_short_branch_names()
    to_delete = [b for b in all_branches if b in closed_pr_branches]

    if not to_delete:
        print("No branches to delete.")
        return

    branch_details = get_branch_details(to_delete)
    print("\nThe following branches will be deleted:")
    for branch in to_delete:
        date = branch_details[branch]
        title = pr_titles[branch]
        print(f"  {branch} {date} - {title}")
    print()

    if args.dry_run:
        print("Dry run - no branches will be deleted")
        return

    response = input("Delete these branches? [y/N] ").lower()
    if response != "y":
        print("Aborted.")
        return

    delete_cmd = ["git", "branch", "-D", *to_delete]
    run_command(git_root, delete_cmd, args)


# Constants for doink installation
DOINK_START_MARKER = "# DOINK REGISTRY-START"
DOINK_END_MARKER = "# DOINK REGISTRY-END"
DOINK_VERSION = "1"


class DoinkSection(NamedTuple):
    """Represents a doink installation section in .zshrc."""

    start_idx: int
    end_idx: int
    content: str
    version: str


def get_doink_section(content: str) -> DoinkSection | None:
    """Find the doink section in .zshrc content.

    Returns a DoinkSection if found, None otherwise.
    """
    if DOINK_START_MARKER not in content:
        return None

    start_idx = content.find(DOINK_START_MARKER)
    end_idx = content.find(DOINK_END_MARKER) + len(DOINK_END_MARKER)
    section_content = content[start_idx:end_idx]

    # Extract version from the start marker
    if match := re.search(rf"{DOINK_START_MARKER} v(\d+)", section_content):
        version = match.group(1)
    else:
        version = "0"  # Default version for old installations

    return DoinkSection(start_idx, end_idx, section_content, version)


def uninstall_command(args: argparse.Namespace) -> None:
    """Remove doink installation from ~/.zshrc."""
    zshrc = Path.home() / ".zshrc"
    if not zshrc.exists():
        print("~/.zshrc not found")
        return

    content = zshrc.read_text()
    section = get_doink_section(content)
    if not section:
        print("No doink installation found in ~/.zshrc")
        return

    if args.dry_run:
        # Create a temporary file with the content that would remain
        import tempfile

        with tempfile.NamedTemporaryFile(mode="w") as tmp:
            new_content = (
                content[: section.start_idx] + content[section.end_idx :].lstrip()
            )
            tmp.write(new_content)
            tmp.flush()
            subprocess.run(["diff", "-u", str(zshrc), tmp.name])
        return

    new_content = content[: section.start_idx] + content[section.end_idx :].lstrip()
    zshrc.write_text(new_content)
    print("Removed doink installation from ~/.zshrc")


def install_command(args: argparse.Namespace) -> None:
    """Install doink shell functions in ~/.zshrc."""
    zshrc = Path.home() / ".zshrc"
    if not zshrc.exists():
        print("~/.zshrc not found")
        return

    content = zshrc.read_text()
    script_path = Path(__file__).resolve()

    # Build the installation content
    install_content = f"""{DOINK_START_MARKER} v{DOINK_VERSION}
DOINK_PATH={script_path}
peep() {{ $DOINK_PATH peep "$@" }}
modified() {{ $DOINK_PATH modified "$@" }}
otherdir() {{ $DOINK_PATH otherdir "$@" }}
meep() {{ $DOINK_PATH meep "$@" }}
prune() {{ $DOINK_PATH prune "$@" }}
{DOINK_END_MARKER}
"""

    # Check if we need to update
    needs_update = False
    section = get_doink_section(content)
    if section:
        if section.version != DOINK_VERSION or str(script_path) not in section.content:
            needs_update = True
            content = content[: section.start_idx] + content[section.end_idx :].lstrip()
        else:
            print("doink is already installed with current version and path")
            return

    # Prepare new content
    new_content = content.rstrip() + "\n\n" + install_content

    if args.dry_run:
        import tempfile

        with tempfile.NamedTemporaryFile(mode="w") as tmp:
            tmp.write(new_content)
            tmp.flush()
            subprocess.run(["diff", "-u", str(zshrc), tmp.name])
        return

    # Write the content
    with zshrc.open("w") as f:
        f.write(new_content)
    print(
        "Installed doink functions in ~/.zshrc"
        if not needs_update
        else "Updated doink installation in ~/.zshrc"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Sentry development utilities")
    subparsers = parser.add_subparsers(dest="command", required=True)

    peep_parser = subparsers.add_parser(
        "peep", help="Find and run tests for modified files"
    )
    peep_parser.add_argument("files", nargs="*", help="Files to test", type=Path)
    peep_parser.add_argument(
        "-o", "--open", action="store_true", help="Open test files in Cursor"
    )
    peep_parser.add_argument("-d", "--dry-run", action="store_true", help="Dry run")

    modified_parser = subparsers.add_parser(
        "modified", help="List modified files (and potentially open them)"
    )
    modified_parser.add_argument(
        "-o", "--open", action="store_true", help="Open files in Cursor"
    )
    modified_parser.add_argument("-t", "--type", choices=("py",), help="File type")
    modified_parser.add_argument("-d", "--dry-run", action="store_true", help="Dry run")

    otherdir_parser = subparsers.add_parser(
        "otherdir", help="Find corresponding test/src directory"
    )
    otherdir_parser.add_argument("-d", "--dry-run", action="store_true", help="Dry run")

    meep_parser = subparsers.add_parser(
        "meep", help="Run mypy on modified Python files"
    )
    meep_parser.add_argument("-d", "--dry-run", action="store_true", help="Dry run")

    prune_parser = subparsers.add_parser(
        "prune",
        help="Delete local branches that correspond to your closed and merged pull requests",
        description="Finds branches from your closed and merged pull requests and deletes them from your local repository. "
        "This helps keep your local repository clean by removing branches that are no longer needed.",
    )
    prune_parser.add_argument(
        "-d",
        "--dry-run",
        action="store_true",
        help="Show what would be deleted without actually deleting",
    )

    install_parser = subparsers.add_parser(
        "install",
        help="Print shell function definitions for direct subcommand usage",
        description="Generates shell function definitions that allow using doink subcommands directly (e.g., 'peep' instead of 'doink peep'). "
        "Add the output to your ~/.zshrc to enable this functionality.",
    )
    install_parser.add_argument(
        "-d",
        "--dry-run",
        action="store_true",
        help="Show what would be done without making changes",
    )

    uninstall_parser = subparsers.add_parser(
        "uninstall",
        help="Remove doink installation from ~/.zshrc",
        description="Removes the doink function definitions from your ~/.zshrc file. "
        "Use -d to see what would be removed without making changes.",
    )
    uninstall_parser.add_argument(
        "-d",
        "--dry-run",
        action="store_true",
        help="Show what would be removed without making changes",
    )

    args = parser.parse_args()

    commands: dict[str, Callable[[argparse.Namespace], None]] = {
        "peep": peep_command,
        "modified": modified_files_command,
        "otherdir": otherdir_command,
        "meep": meep_command,
        "prune": prune_command,
        "install": install_command,
        "uninstall": uninstall_command,
    }

    commands[args.command](args)


if __name__ == "__main__":
    main()