.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     # If .git is a symlink, worktree result might be the target.
 143     if worktree == str(git_dir().resolve()):
 144         return True
 145
 146     # NB: worktree path is supposed to be absolute from for-each-ref, but it's
 147     # not always, so we have to resolve it.  https://crbug.com/git/88
 148     worktree = (git_dir() / worktree).resolve()
 149     return worktree == top_dir()
 150
 151
 152 class AppendOption(argparse.Action):
 153     """Append the command line option (with no arguments) to dest.
 154
 155     parser.add_argument('-b', '--barg', dest='out', action='append_option')
 156     options = parser.parse_args(['-b', '--barg'])
 157     options.out == ['-b', '--barg']
 158     """
 159
 160     def __init__(self, option_strings, dest, **kwargs):
 161         if 'nargs' in kwargs:
 162             raise ValueError('nargs is not supported for append_option action')
 163         super().__init__(option_strings, dest, nargs=0, **kwargs)
 164
 165     def __call__(self, parser, namespace, values, option_string=None):
 166         if getattr(namespace, self.dest, None) is None:
 167             setattr(namespace, self.dest, [])
 168         getattr(namespace, self.dest).append(option_string)
 169
 170
 171 def get_parser():
 172     """Get CLI parser."""
 173     parser = argparse.ArgumentParser(description=__doc__)
 174     parser.add_argument(
 175         '--catchup', action='store_true',
 176         help='run git-rb-catchup when rebasing')
 177     parser.add_argument(
 178         '-d', '--debug', action='store_true',
 179         help='enable debug output')
 180     parser.add_argument(
 181         '-q', '--quiet', dest='git_options', action=AppendOption, default=['-q'],
 182         help='passthru to git rebase')
 183     return parser
 184
 185
 186 def main(argv):
 187     """The main entry point for scripts."""
 188     parser = get_parser()
 189     opts = parser.parse_args(argv)
 190
 191     global DEBUG
 192     DEBUG = opts.debug
 193
 194     # Switch to the top dir in case the working dir doesn't exist in every branch.
 195     try:
 196         topdir = top_dir()
 197     except subprocess.CalledProcessError as e:
 198         sys.exit(f'{os.path.basename(sys.argv[0])}: {Path.cwd()}:\n{e}\n{e.stderr.strip()}')
 199     os.chdir(topdir)
 200
 201     # Example output:
 202     #  ||m|refs/remotes/origin/master|ahead 2, behind 203
 203     # *||master|refs/remotes/origin/master|ahead 1
 204     #  |/usr/local/src/gnu/gdb/build/release|release||behind 10
 205     #  ||s-stash||
 206     state = git(
 207         ['for-each-ref', '--format=%(HEAD)|%(worktreepath)|%(refname:short)|%(upstream)|%(upstream:track,nobracket)',
 208          'HEAD', 'refs/heads/*']).stdout.splitlines()
 209
 210     curr_state = None
 211     branch_width = 0
 212     local_count = 0
 213     for line in state:
 214         head, worktreepath, branch, tracking, ahead_behind = line.split('|')
 215         branch_width = max(branch_width, len(branch))
 216         if head == '*':
 217             curr_state = branch
 218             local_count += 1
 219         elif worktree_is_local(worktreepath):
 220             local_count += 1
 221         else:
 222             dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
 223     # If there are no branches to rebase, go silently.
 224     if not local_count:
 225         return 0
 226     if not curr_state:
 227         # Are we in a detached head state?
 228         if not git(['symbolic-ref', '-q', 'HEAD'], check=False).returncode:
 229             fatal('unable to resolve current branch')
 230         curr_state = git(['rev-parse', 'HEAD']).stdout.strip()
 231
 232     switched_head = False
 233     branches = {}
 234     for line in state:
 235         _, worktreepath, branch, tracking, ahead_behind = line.split('|')
 236
 237         # If it's a branch in another worktree, ignore it.
 238         if not worktree_is_local(worktreepath):
 239             dbg(f'{worktreepath}:{branch}: Skipping branch checked out in diff worktree')
 240             continue
 241
 242         print(f'{Color.BRACKET}### {Color.GOOD}{branch:{branch_width}}{Color.NORMAL} ',
 243               end='', flush=True)
 244         if not tracking:
 245             print(f'{Color.WARN}skipping{Color.NORMAL} due to missing merge branch')
 246             continue
 247
 248         m = re.match(r'ahead ([0-9]+)', ahead_behind)
 249         ahead = int(m.group(1)) if m else 0
 250         m = re.search(r'behind ([0-9]+)', ahead_behind)
 251         behind = int(m.group(1)) if m else 0
 252         if not behind:
 253             print(Color.good('up-to-date'))
 254             continue
 255         elif not ahead:
 256             # If we haven't switched the checkout, update-ref on current HEAD
 257             # will get us into a dirty checkout, so use merge to handle it.
 258             if switched_head or curr_state != branch:
 259                 git(['update-ref', f'refs/heads/{branch}', tracking])
 260                 print('fast forwarded [updated ref]')
 261             else:
 262                 result = git(['merge', '-q', '--ff-only'], check=False)
 263                 if result.returncode:
 264                     print(Color.bad('unable to merge') + '\n' + result.stdout.strip())
 265                 else:
 266                     print('fast forwarded [merged]')
 267             continue
 268
 269         # Skip this ref if tree is in a bad state.
 270         if rebase_inprogress():
 271             print(Color.bad('skipping due to active rebase'))
 272             continue
 273         if cherry_pick_inprogress():
 274             print(Color.bad('skipping due to active cherry-pick'))
 275             continue
 276         if checkout_is_dirty():
 277             print(Color.bad('unable to rebase: tree is dirty'))
 278             continue
 279
 280         print(f'rebasing [{ahead_behind}] ', end='', flush=True)
 281         result = git(['checkout', '-q', branch], check=False)
 282         if result.returncode:
 283             print(Color.bad('unable to checkout') + '\n' + result.stdout.strip())
 284             continue
 285         switched_head = True
 286         if opts.catchup:
 287             print()
 288             result = git(['rb-catchup'], capture_output=False, check=False)
 289         else:
 290             result = git(['rebase'] + opts.git_options, check=False)
 291             if result.returncode:
 292                 git(['rebase', '--abort'], check=False)
 293                 print(Color.bad('failed') + '\n' + result.stdout.strip())
 294             else:
 295                 print(Color.good('OK'))
 296
 297     if switched_head:
 298         git(['checkout', '-q', curr_state])
 299     return 0
 300
 301
 302 if __name__ == '__main__':
 303     sys.exit(main(sys.argv[1:]))