git-rb-all: fix update-ref usage with current HEAD, and make nops faster

[home.git] / .bin / git-rb-all

#!/usr/bin/env python3

"""Helper to rebase all local branches."""

import argparse
import functools
import os
from pathlib import Path
import re
import subprocess
import sys


# Not sure if newer features are used.
assert sys.version_info >= (3, 6), f'Python 3.6+ required but found {sys.version}'


PROG = os.path.basename(__file__)


class Terminal:
    """Terminal escape sequences."""

    CSI_PREFIX = '\033['
    SGR_SUFFIX = 'm'
    NORMAL = ''
    BOLD = '1'
    _BLACK, _RED, _GREEN, _YELLOW, _BLUE, _MAGENTA, _CYAN, _WHITE = range(0, 8)
    _FG = 30
    _BG = 40
    FG_BLACK = str(_FG + _BLACK)
    FG_RED = str(_FG + _RED)
    FG_GREEN = str(_FG + _GREEN)
    FG_YELLOW = str(_FG + _YELLOW)
    FG_BLUE = str(_FG + _BLUE)
    FG_MAGENTA = str(_FG + _MAGENTA)
    FG_CYAN = str(_FG + _CYAN)
    FG_WHITE = str(_FG + _WHITE)


class Color:
    """Helper colors."""

    _combine = lambda *args: Terminal.CSI_PREFIX + ';'.join(args) + Terminal.SGR_SUFFIX
    NORMAL = _combine(Terminal.NORMAL)
    GOOD = _combine(Terminal.FG_GREEN)
    WARN = _combine(Terminal.FG_YELLOW)
    BAD = _combine(Terminal.FG_RED)
    HILITE = _combine(Terminal.BOLD, Terminal.FG_CYAN)
    BRACKET = _combine(Terminal.BOLD, Terminal.FG_BLUE)


def fatal(msg):
    """Show an error |msg| then exit."""
    print(f'{Color.BAD}{PROG}: error: {msg}{Color.NORMAL}', file=sys.stderr)
    sys.exit(1)


def git(args, **kwargs):
    """Run git."""
    kwargs.setdefault('check', True)
    kwargs.setdefault('stdout', subprocess.PIPE)
    kwargs.setdefault('stderr', subprocess.STDOUT)
    kwargs.setdefault('encoding', 'utf-8')
    #print('+', 'git', *args)
    return subprocess.run(['git'] + args, **kwargs)


@functools.lru_cache(maxsize=None)
def checkout_is_dirty():
    """Determine whether the checkout is dirty (e.g. modified files)."""
    output = git(['status', '--porcelain']).stdout
    return any(x for x in output.splitlines() if '?' not in x[0:2])


@functools.lru_cache(maxsize=None)
def rebase_inprogress():
    """Determine whether a rebase is already in progress."""
    output = git(['rev-parse', '--git-path', 'rebase-merge']).stdout.strip()
    if Path(output).exists():
        return True

    output = git(['rev-parse', '--git-path', 'rebase-apply']).stdout.strip()
    if Path(output).exists():
        return True

    return False


class AppendOption(argparse.Action):
    """Append the command line option (with no arguments) to dest.

    parser.add_argument('-b', '--barg', dest='out', action='append_option')
    options = parser.parse_args(['-b', '--barg'])
    options.out == ['-b', '--barg']
    """

    def __init__(self, option_strings, dest, **kwargs):
        if 'nargs' in kwargs:
            raise ValueError('nargs is not supported for append_option action')
        super().__init__(option_strings, dest, nargs=0, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        if getattr(namespace, self.dest, None) is None:
            setattr(namespace, self.dest, [])
        getattr(namespace, self.dest).append(option_string)


def get_parser():
    """Get CLI parser."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        '--catchup', action='store_true',
        help='run git-rb-catchup when rebasing')
    parser.add_argument(
        '-q', '--quiet', dest='git_options', action=AppendOption, default=['-q'],
        help='passthru to git rebase')
    return parser


def main(argv):
    """The main entry point for scripts."""
    parser = get_parser()
    opts = parser.parse_args(argv)

    # Skip if rebase is in progress.
    if rebase_inprogress():
        fatal('skipping due to active rebase')

    # Switch to the top dir in case the working dir doesn't exist in every branch.
    topdir = git(['rev-parse', '--show-toplevel']).stdout.strip()
    os.chdir(topdir)

    # Example output:
    #  ||m|refs/remotes/origin/master|ahead 2, behind 203
    # *||master|refs/remotes/origin/master|ahead 1
    #  |/usr/local/src/gnu/gdb/build/release|release||behind 10
    #  ||s-stash||
    state = git(
        ['for-each-ref', '--format=%(HEAD)|%(worktreepath)|%(refname:short)|%(upstream)|%(upstream:track,nobracket)',
         'refs/heads/*']).stdout.splitlines()

    curr_state = None
    branch_width = 0
    local_count = 0
    for line in state:
        head, worktreepath, branch, tracking, ahead_behind = line.split('|')
        branch_width = max(branch_width, len(branch))
        if head == '*':
            curr_state = branch
        if not (worktreepath and worktreepath != topdir):
            local_count += 1
    # If there are no branches to rebase, go silently.
    if not local_count:
        return 0
    if not curr_state:
        fatal('unable to resolve current branch')

    switched_head = False
    branches = {}
    for line in state:
        _, worktreepath, branch, tracking, ahead_behind = line.split('|')

        # If it's a branch in another worktree, ignore it.
        if worktreepath and worktreepath != topdir:
            continue

        print(f'{Color.BRACKET}### {Color.GOOD}{branch:{branch_width}}{Color.NORMAL} ',
              end='', flush=True)
        if not tracking:
            print(f'{Color.WARN}skipping{Color.NORMAL} due to missing merge branch')
            continue

        m = re.match(r'ahead ([0-9]+)', ahead_behind)
        ahead = int(m.group(1)) if m else 0
        m = re.search(r'behind ([0-9]+)', ahead_behind)
        behind = int(m.group(1)) if m else 0
        if not behind:
            print('Up-to-date!')
            continue
        elif not ahead:
            # If we haven't switched the checkout, update-ref on current HEAD
            # will get us into a dirty checkout, so use merge to handle it.
            if switched_head or curr_state != branch:
                git(['update-ref', f'refs/heads/{branch}', tracking])
                print('fast forwarded [updated ref]')
            else:
                result = git(['merge', '-q', '--ff-only'], check=False)
                if result.returncode:
                    print(f'{Color.BAD}unable to merge{Color.NORMAL}\n' + result.stdout.strip())
                else:
                    print('fast forwarded [merged]')
            continue

        if checkout_is_dirty():
            print(f'{Color.BAD}unable to rebase: tree is dirty{Color.NORMAL}')
            continue

        print(f'rebasing [{ahead_behind}] ', end='', flush=True)
        git(['checkout', '-q', branch])
        switched_head = True
        if opts.catchup:
            print()
            result = git(['rb-catchup'], capture_output=False, check=False)
        else:
            result = git(['rebase'] + opts.git_options, check=False)
            if result.returncode:
                git(['rebase', '--abort'])
                print(f'{Color.BAD}failed{Color.NORMAL}\n' + result.stdout.strip())
            else:
                print('OK!')

    if switched_head:
        git(['checkout', '-q', curr_state])
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))