.bin/git-rb-all

   1 #!/usr/bin/env python3
   2
   3 """Helper to rebase all local branches."""
   4
   5 import argparse
   6 import functools
   7 import os
   8 from pathlib import Path
   9 import re
  10 import shlex
  11 import subprocess
  12 import sys
  13
  14
  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
  19 PROG = os.path.basename(__file__)
  20 DEBUG = False
  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
  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
  62
  63 def dbg(*args, **kwargs):
  64     """Print a debug |msg|."""
  65     if DEBUG:
  66         print(*args, file=sys.stderr, **kwargs)
  67
  68
  69 def fatal(msg):
  70     """Show an error |msg| then exit."""
  71     print(Color.bad(f'{PROG}: error: {msg}'), file=sys.stderr)
  72     sys.exit(1)
  73
  74
  75 def git(args, **kwargs):
  76     """Run git."""
  77     kwargs.setdefault('check', True)
  78     kwargs.setdefault('stdout', subprocess.PIPE)
  79     kwargs.setdefault('stderr', subprocess.STDOUT)
  80     kwargs.setdefault('encoding', 'utf-8')
  81     if DEBUG:
  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
  92
  93
  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)
 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
 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
 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
 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')
 173     parser.add_argument(
 174         '-d', '--debug', action='store_true',
 175         help='enable debug output')
 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
 187     global DEBUG
 188     DEBUG = opts.debug
 189
 190     # Switch to the top dir in case the working dir doesn't exist in every branch.
 191     try:
 192         topdir = top_dir()
 193     except subprocess.CalledProcessError as e:
 194         sys.exit(f'{os.path.basename(sys.argv[0])}: {Path.cwd()}:\n{e}\n{e.stderr.strip()}')
 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)',
 204          'HEAD', 'refs/heads/*']).stdout.splitlines()
 205
 206     curr_state = None
 207     branch_width = 0
 208     local_count = 0
 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
 214             local_count += 1
 215         elif worktree_is_local(worktreepath):
 216             local_count += 1
 217         else:
 218             dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
 219     # If there are no branches to rebase, go silently.
 220     if not local_count:
 221         return 0
 222     if not curr_state:
 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()
 227
 228     switched_head = False
 229     branches = {}
 230     for line in state:
 231         _, worktreepath, branch, tracking, ahead_behind = line.split('|')
 232
 233         # If it's a branch in another worktree, ignore it.
 234         if not worktree_is_local(worktreepath):
 235             dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
 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:
 249             print(Color.good('up-to-date'))
 250             continue
 251         elif not ahead:
 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:
 260                     print(Color.bad('unable to merge') + '\n' + result.stdout.strip())
 261                 else:
 262                     print('fast forwarded [merged]')
 263             continue
 264
 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
 272         if checkout_is_dirty():
 273             print(Color.bad('unable to rebase: tree is dirty'))
 274             continue
 275
 276         print(f'rebasing [{ahead_behind}] ', end='', flush=True)
 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
 281         switched_head = True
 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:
 288                 git(['rebase', '--abort'], check=False)
 289                 print(Color.bad('failed') + '\n' + result.stdout.strip())
 290             else:
 291                 print(Color.good('OK'))
 292
 293     if switched_head:
 294         git(['checkout', '-q', curr_state])
 295     return 0
 296
 297
 298 if __name__ == '__main__':
 299     sys.exit(main(sys.argv[1:]))