[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 shlex
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__)
DEBUG = False


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)

    @classmethod
    def good(cls, msg):
        return cls.GOOD + msg + cls.NORMAL

    @classmethod
    def bad(cls, msg):
        return cls.BAD + msg + cls.NORMAL


def dbg(*args, **kwargs):
    """Print a debug |msg|."""
    if DEBUG:
        print(*args, file=sys.stderr, **kwargs)


def fatal(msg):
    """Show an error |msg| then exit."""
    print(Color.bad(f'{PROG}: error: {msg}'), 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')
    if DEBUG:
        dbg('+', 'git', shlex.join(args))
    ret = subprocess.run(['git'] + args, **kwargs)
    if DEBUG:
        if ret.stdout:
            dbg(ret.stdout.rstrip())
        if ret.stderr:
            dbg('stderr =', ret.stderr)
        if ret.returncode:
            dbg('++ exit', ret.returncode)
    return ret


@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


@functools.lru_cache(maxsize=None)
def cherry_pick_inprogress():
    """Determine whether a cherry-pick is in progress."""
    output = git(['rev-parse', '--git-path', 'CHERRY_PICK_HEAD']).stdout.strip()
    return Path(output).exists()


@functools.lru_cache(maxsize=None)
def top_dir() -> Path:
    """Find the top dir of the git checkout."""
    output = git(['rev-parse', '--show-toplevel'], stderr=subprocess.PIPE).stdout.strip()
    return Path(output).resolve()


@functools.lru_cache(maxsize=None)
def git_dir() -> Path:
    """Find the internal git dir for this project."""
    output = git(['rev-parse', '--git-dir']).stdout.strip()
    return Path(output).resolve()


@functools.lru_cache(maxsize=None)
def worktree_is_local(worktree: str) -> bool:
    """See whether |worktree| is the cwd git repo."""
    if not worktree:
        return True

    # NB: worktree path is supposed to be absolute from for-each-ref, but it's
    # not always, so we have to resolve it.  https://crbug.com/git/88
    worktree = (git_dir() / worktree).resolve()
    return worktree == top_dir()


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(
        '-d', '--debug', action='store_true',
        help='enable debug output')
    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)

    global DEBUG
    DEBUG = opts.debug

    # Switch to the top dir in case the working dir doesn't exist in every branch.
    try:
        topdir = top_dir()
    except subprocess.CalledProcessError as e:
        sys.exit(f'{os.path.basename(sys.argv[0])}: {Path.cwd()}:\n{e}\n{e.stderr.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)',
         'HEAD', '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
            local_count += 1
        elif worktree_is_local(worktreepath):
            local_count += 1
        else:
            dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
    # If there are no branches to rebase, go silently.
    if not local_count:
        return 0
    if not curr_state:
        # Are we in a detached head state?
        if not git(['symbolic-ref', '-q', 'HEAD'], check=False).returncode:
            fatal('unable to resolve current branch')
        curr_state = git(['rev-parse', 'HEAD']).stdout.strip()

    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 not worktree_is_local(worktreepath):
            dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
            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(Color.good('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(Color.bad('unable to merge') + '\n' + result.stdout.strip())
                else:
                    print('fast forwarded [merged]')
            continue

        # Skip this ref if tree is in a bad state.
        if rebase_inprogress():
            print(Color.bad('skipping due to active rebase'))
            continue
        if cherry_pick_inprogress():
            print(Color.bad('skipping due to active cherry-pick'))
            continue
        if checkout_is_dirty():
            print(Color.bad('unable to rebase: tree is dirty'))
            continue

        print(f'rebasing [{ahead_behind}] ', end='', flush=True)
        result = git(['checkout', '-q', branch], check=False)
        if result.returncode:
            print(Color.bad('unable to checkout') + '\n' + result.stdout.strip())
            continue
        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'], check=False)
                print(Color.bad('failed') + '\n' + result.stdout.strip())
            else:
                print(Color.good('OK'))

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


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))
Commit	Line	Data
2ea871e7 MF	1	#!/usr/bin/env python3
	2
	3	"""Helper to rebase all local branches."""
	4
	5	import argparse
eb86c710	6	import functools
2ea871e7 MF	7	import os
	8	from pathlib import Path
	9	import re
8c438e9f	10	import shlex
2ea871e7 MF	11	import subprocess
	12	import sys
	13
	14
eb86c710 MF	15	# Not sure if newer features are used.
	16	assert sys.version_info >= (3, 6), f'Python 3.6+ required but found {sys.version}'
	17
	18
2ea871e7	19	PROG = os.path.basename(__file__)
669b9886	20	DEBUG = False
2ea871e7 MF	21
	22
	23	class Terminal:
	24	"""Terminal escape sequences."""
	25
	26	CSI_PREFIX = '\033['
	27	SGR_SUFFIX = 'm'
	28	NORMAL = ''
	29	BOLD = '1'
	30	_BLACK, _RED, _GREEN, _YELLOW, _BLUE, _MAGENTA, _CYAN, _WHITE = range(0, 8)
	31	_FG = 30
	32	_BG = 40
	33	FG_BLACK = str(_FG + _BLACK)
	34	FG_RED = str(_FG + _RED)
	35	FG_GREEN = str(_FG + _GREEN)
	36	FG_YELLOW = str(_FG + _YELLOW)
	37	FG_BLUE = str(_FG + _BLUE)
	38	FG_MAGENTA = str(_FG + _MAGENTA)
	39	FG_CYAN = str(_FG + _CYAN)
	40	FG_WHITE = str(_FG + _WHITE)
	41
	42
	43	class Color:
	44	"""Helper colors."""
	45
	46	_combine = lambda *args: Terminal.CSI_PREFIX + ';'.join(args) + Terminal.SGR_SUFFIX
	47	NORMAL = _combine(Terminal.NORMAL)
	48	GOOD = _combine(Terminal.FG_GREEN)
	49	WARN = _combine(Terminal.FG_YELLOW)
	50	BAD = _combine(Terminal.FG_RED)
	51	HILITE = _combine(Terminal.BOLD, Terminal.FG_CYAN)
	52	BRACKET = _combine(Terminal.BOLD, Terminal.FG_BLUE)
	53
fb5d2cb6 MF	54	@classmethod
	55	def good(cls, msg):
	56	return cls.GOOD + msg + cls.NORMAL
	57
	58	@classmethod
	59	def bad(cls, msg):
	60	return cls.BAD + msg + cls.NORMAL
	61
2ea871e7	62
8c438e9f	63	def dbg(args, *kwargs):
669b9886 MF	64	"""Print a debug \|msg\|."""
669b9886 MF	65	if DEBUG:
8c438e9f	66	print(args, file=sys.stderr, *kwargs)
669b9886 MF	67
669b9886 MF	68
65a6ea84 MF	69	def fatal(msg):
65a6ea84 MF	70	"""Show an error \|msg\| then exit."""
fb5d2cb6	71	print(Color.bad(f'{PROG}: error: {msg}'), file=sys.stderr)
65a6ea84 MF	72	sys.exit(1)
	73
	74
2ea871e7 MF	75	def git(args, **kwargs):
	76	"""Run git."""
	77	kwargs.setdefault('check', True)
eb86c710 MF	78	kwargs.setdefault('stdout', subprocess.PIPE)
eb86c710 MF	79	kwargs.setdefault('stderr', subprocess.STDOUT)
2ea871e7	80	kwargs.setdefault('encoding', 'utf-8')
669b9886	81	if DEBUG:
8c438e9f MF	82	dbg('+', 'git', shlex.join(args))
	83	ret = subprocess.run(['git'] + args, **kwargs)
	84	if DEBUG:
	85	if ret.stdout:
	86	dbg(ret.stdout.rstrip())
	87	if ret.stderr:
	88	dbg('stderr =', ret.stderr)
	89	if ret.returncode:
	90	dbg('++ exit', ret.returncode)
	91	return ret
2ea871e7 MF	92
2ea871e7 MF	93
eb86c710 MF	94	@functools.lru_cache(maxsize=None)
	95	def checkout_is_dirty():
	96	"""Determine whether the checkout is dirty (e.g. modified files)."""
	97	output = git(['status', '--porcelain']).stdout
	98	return any(x for x in output.splitlines() if '?' not in x[0:2])
	99
	100
	101	@functools.lru_cache(maxsize=None)
2ea871e7 MF	102	def rebase_inprogress():
	103	"""Determine whether a rebase is already in progress."""
	104	output = git(['rev-parse', '--git-path', 'rebase-merge']).stdout.strip()
	105	if Path(output).exists():
	106	return True
	107
	108	output = git(['rev-parse', '--git-path', 'rebase-apply']).stdout.strip()
	109	if Path(output).exists():
	110	return True
	111
	112	return False
	113
	114
fb5d2cb6 MF	115	@functools.lru_cache(maxsize=None)
	116	def cherry_pick_inprogress():
	117	"""Determine whether a cherry-pick is in progress."""
	118	output = git(['rev-parse', '--git-path', 'CHERRY_PICK_HEAD']).stdout.strip()
	119	return Path(output).exists()
	120
	121
8c438e9f MF	122	@functools.lru_cache(maxsize=None)
	123	def top_dir() -> Path:
	124	"""Find the top dir of the git checkout."""
	125	output = git(['rev-parse', '--show-toplevel'], stderr=subprocess.PIPE).stdout.strip()
	126	return Path(output).resolve()
	127
	128
	129	@functools.lru_cache(maxsize=None)
	130	def git_dir() -> Path:
	131	"""Find the internal git dir for this project."""
	132	output = git(['rev-parse', '--git-dir']).stdout.strip()
	133	return Path(output).resolve()
	134
	135
	136	@functools.lru_cache(maxsize=None)
	137	def worktree_is_local(worktree: str) -> bool:
	138	"""See whether \|worktree\| is the cwd git repo."""
	139	if not worktree:
	140	return True
	141
	142	# NB: worktree path is supposed to be absolute from for-each-ref, but it's
	143	# not always, so we have to resolve it. https://crbug.com/git/88
	144	worktree = (git_dir() / worktree).resolve()
	145	return worktree == top_dir()
	146
	147
2ea871e7 MF	148	class AppendOption(argparse.Action):
	149	"""Append the command line option (with no arguments) to dest.
	150
	151	parser.add_argument('-b', '--barg', dest='out', action='append_option')
	152	options = parser.parse_args(['-b', '--barg'])
	153	options.out == ['-b', '--barg']
	154	"""
	155
	156	def __init__(self, option_strings, dest, **kwargs):
	157	if 'nargs' in kwargs:
	158	raise ValueError('nargs is not supported for append_option action')
	159	super().__init__(option_strings, dest, nargs=0, **kwargs)
	160
	161	def __call__(self, parser, namespace, values, option_string=None):
	162	if getattr(namespace, self.dest, None) is None:
	163	setattr(namespace, self.dest, [])
	164	getattr(namespace, self.dest).append(option_string)
	165
	166
	167	def get_parser():
	168	"""Get CLI parser."""
	169	parser = argparse.ArgumentParser(description=__doc__)
	170	parser.add_argument(
	171	'--catchup', action='store_true',
	172	help='run git-rb-catchup when rebasing')
669b9886 MF	173	parser.add_argument(
	174	'-d', '--debug', action='store_true',
	175	help='enable debug output')
2ea871e7 MF	176	parser.add_argument(
	177	'-q', '--quiet', dest='git_options', action=AppendOption, default=['-q'],
	178	help='passthru to git rebase')
	179	return parser
	180
	181
	182	def main(argv):
	183	"""The main entry point for scripts."""
	184	parser = get_parser()
	185	opts = parser.parse_args(argv)
	186
669b9886 MF	187	global DEBUG
	188	DEBUG = opts.debug
	189
2ea871e7	190	# Switch to the top dir in case the working dir doesn't exist in every branch.
669b9886	191	try:
8c438e9f	192	topdir = top_dir()
669b9886 MF	193	except subprocess.CalledProcessError as e:
669b9886 MF	194	sys.exit(f'{os.path.basename(sys.argv[0])}: {Path.cwd()}:\n{e}\n{e.stderr.strip()}')
2ea871e7 MF	195	os.chdir(topdir)
	196
	197	# Example output:
	198	# \|\|m\|refs/remotes/origin/master\|ahead 2, behind 203
	199	# *\|\|master\|refs/remotes/origin/master\|ahead 1
	200	# \|/usr/local/src/gnu/gdb/build/release\|release\|\|behind 10
	201	# \|\|s-stash\|\|
	202	state = git(
	203	['for-each-ref', '--format=%(HEAD)\|%(worktreepath)\|%(refname:short)\|%(upstream)\|%(upstream:track,nobracket)',
fb5d2cb6	204	'HEAD', 'refs/heads/*']).stdout.splitlines()
2ea871e7 MF	205
	206	curr_state = None
	207	branch_width = 0
65a6ea84	208	local_count = 0
2ea871e7 MF	209	for line in state:
	210	head, worktreepath, branch, tracking, ahead_behind = line.split('\|')
	211	branch_width = max(branch_width, len(branch))
	212	if head == '*':
	213	curr_state = branch
65a6ea84	214	local_count += 1
8c438e9f	215	elif worktree_is_local(worktreepath):
669b9886 MF	216	local_count += 1
	217	else:
	218	dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
65a6ea84 MF	219	# If there are no branches to rebase, go silently.
	220	if not local_count:
	221	return 0
2ea871e7	222	if not curr_state:
fb5d2cb6 MF	223	# Are we in a detached head state?
	224	if not git(['symbolic-ref', '-q', 'HEAD'], check=False).returncode:
	225	fatal('unable to resolve current branch')
	226	curr_state = git(['rev-parse', 'HEAD']).stdout.strip()
2ea871e7	227
eb86c710	228	switched_head = False
2ea871e7 MF	229	branches = {}
2ea871e7 MF	230	for line in state:
eb86c710	231	_, worktreepath, branch, tracking, ahead_behind = line.split('\|')
2ea871e7 MF	232
2ea871e7 MF	233	# If it's a branch in another worktree, ignore it.
8c438e9f MF	234	if not worktree_is_local(worktreepath):
8c438e9f MF	235	dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
2ea871e7 MF	236	continue
	237
	238	print(f'{Color.BRACKET}### {Color.GOOD}{branch:{branch_width}}{Color.NORMAL} ',
	239	end='', flush=True)
	240	if not tracking:
	241	print(f'{Color.WARN}skipping{Color.NORMAL} due to missing merge branch')
	242	continue
	243
	244	m = re.match(r'ahead ([0-9]+)', ahead_behind)
	245	ahead = int(m.group(1)) if m else 0
	246	m = re.search(r'behind ([0-9]+)', ahead_behind)
	247	behind = int(m.group(1)) if m else 0
	248	if not behind:
fb5d2cb6	249	print(Color.good('up-to-date'))
2ea871e7 MF	250	continue
2ea871e7 MF	251	elif not ahead:
eb86c710 MF	252	# If we haven't switched the checkout, update-ref on current HEAD
	253	# will get us into a dirty checkout, so use merge to handle it.
	254	if switched_head or curr_state != branch:
	255	git(['update-ref', f'refs/heads/{branch}', tracking])
	256	print('fast forwarded [updated ref]')
	257	else:
	258	result = git(['merge', '-q', '--ff-only'], check=False)
	259	if result.returncode:
fb5d2cb6	260	print(Color.bad('unable to merge') + '\n' + result.stdout.strip())
eb86c710 MF	261	else:
eb86c710 MF	262	print('fast forwarded [merged]')
2ea871e7 MF	263	continue
2ea871e7 MF	264
fb5d2cb6 MF	265	# Skip this ref if tree is in a bad state.
	266	if rebase_inprogress():
	267	print(Color.bad('skipping due to active rebase'))
	268	continue
	269	if cherry_pick_inprogress():
	270	print(Color.bad('skipping due to active cherry-pick'))
	271	continue
eb86c710	272	if checkout_is_dirty():
fb5d2cb6	273	print(Color.bad('unable to rebase: tree is dirty'))
65a6ea84 MF	274	continue
65a6ea84 MF	275
2ea871e7	276	print(f'rebasing [{ahead_behind}] ', end='', flush=True)
fb5d2cb6 MF	277	result = git(['checkout', '-q', branch], check=False)
	278	if result.returncode:
	279	print(Color.bad('unable to checkout') + '\n' + result.stdout.strip())
	280	continue
eb86c710	281	switched_head = True
2ea871e7 MF	282	if opts.catchup:
	283	print()
	284	result = git(['rb-catchup'], capture_output=False, check=False)
	285	else:
	286	result = git(['rebase'] + opts.git_options, check=False)
	287	if result.returncode:
8c438e9f	288	git(['rebase', '--abort'], check=False)
fb5d2cb6	289	print(Color.bad('failed') + '\n' + result.stdout.strip())
2ea871e7	290	else:
fb5d2cb6	291	print(Color.good('OK'))
2ea871e7	292
eb86c710 MF	293	if switched_head:
eb86c710 MF	294	git(['checkout', '-q', curr_state])
2ea871e7 MF	295	return 0
	296
	297
	298	if __name__ == '__main__':
	299	sys.exit(main(sys.argv[1:]))