Clean git branches and stale references

git_clean_branch.py

#!/usr/bin/env python
import sys
import argparse
import subprocess


REMOTE = 'origin'


class CmdError(RuntimeError):
    pass


def run_cmd(cmd, dry_run=False, log=True, print_stderr=True, on_error=''):
    # runs command, captures stdout, and return output lines
    # if command exits with non-zero status raises a CmdError
    if log or dry_run:
        print(('(dry-run) ' if dry_run else '-> ') + ' '.join(cmd))
    if dry_run:
        return []
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=None if print_stderr else subprocess.PIPE)
    out, _ = proc.communicate()
    if proc.returncode == 0:
        out = out.decode('utf-8').strip()
        return out.split('\n') if out else []
    else:
        raise CmdError(on_error or ' '.join(cmd))


def get_git_commit_for(commitish):
    # returns the commit for a commitish: branch, tag, etc.
    cmd = ['git', 'rev-parse', commitish]
    return run_cmd(cmd, log=False, on_error='unrecognized commitish ' + commitish)[0]


def get_remote_tracking_branch_for(branch):
    # returns the remote tracking branch or None if none.
    cmd = ['git', 'rev-parse', '--abbrev-ref', branch + '@{upstream}']
    try:
        for line in run_cmd(cmd, log=False, print_stderr=False):
            return line
    except CmdError:
        pass


def verify_master_status():
    # asserts that master and origin/master are at the same commit
    # assumes git fetch has already been called
    local = get_git_commit_for('master')
    remote = get_git_commit_for('origin/master')
    if local != remote:
        raise CmdError('master and origin/master are not in the same place: %s vs %s' %
                       (local[:8], remote[:8]))


def git_fetch():
    run_cmd(['git', 'fetch'])


def get_current_branch():
    return run_cmd(['git', 'rev-parse', '--abbrev-ref', 'HEAD'], log=False)[0]


def branch_is_protected(branch):
    return 'HEAD' in branch or 'master' in branch or 'release' in branch


def get_merged_local_branches():
    # yields git branches, excluding current and protected branches, that have
    # been merged into <remote>/master.
    # NOTE assumes git fetch has already been called.
    current_branch = get_current_branch()
    cmd = [
        'git', 'branch',
        '--format', '%(refname:short)', # drop formatting, including the asterisk
        '--merged', 'origin/master'     # only include merged branches
    ]
    for line in run_cmd(cmd, log=False):
        branch = line.strip()
        if branch != current_branch and not branch_is_protected(branch):
            yield branch


def get_pushed_local_branches():
    # yields git branches, excluding current and protected branches, that are
    # up to date with remote.
    # NOTE assumes git fetch has been called.
    current_branch = get_current_branch()
    cmd = [
        'git', 'branch',
        '--format', '%(refname:short)', # drop formatting, including the asterisk
    ]
    for line in run_cmd(cmd, log=False):
        branch = line.strip()
        if branch == current_branch or branch_is_protected(branch):
            continue

        rt_branch = get_remote_tracking_branch_for(branch)
        if rt_branch is None:
            continue

        local = get_git_commit_for(branch)
        remote = get_git_commit_for(rt_branch)
        if local == remote:
            yield branch


def delete_local_branch(branch, dry_run=False):
    cmd = ['git', 'branch', '-d', branch]
    err = 'Failed to delete branch ' + branch
    run_cmd(cmd, dry_run=dry_run, on_error=err)


def delete_merged_local_branches(dry_run=False):
    for branch in get_merged_local_branches():
        delete_local_branch(branch, dry_run=dry_run)


def delete_stale_remote_tracking_branches(dry_run=False):
    # removes local tracking branches for branches that have been removed from
    # remote (regardless of merge status)
    for line in run_cmd(['git', 'remote', 'prune', 'origin'], dry_run=dry_run):
        print(line)


def delete_pushed_local_branches(dry_run=False):
    for branch in get_pushed_local_branches():
        delete_local_branch(branch, dry_run=dry_run)


def clean_branches(dry_run):
    git_fetch()
    verify_master_status()
    delete_merged_local_branches(dry_run=dry_run)
    delete_stale_remote_tracking_branches(dry_run=dry_run)
    delete_pushed_local_branches(dry_run=dry_run)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="""
        Clean git branches and stale references.

        - remove all local branches that are merged in {remote}/master
        - remove stale remote tracking branches: those  that are no longer on
          {remote}.
        - remove all local branches that are up to date with their upstream in
          remote.

        Notes:

        - Nothing is pushed to / modified in remote.
        - Nothing is every done to protected branches (master and releases)
        - Nothing is done to the current checked out branch.
    """.format(remote=REMOTE))
    parser.add_argument('--dry-run', action='store_true',
                        help='print git commands instead of running them')
    args = parser.parse_args()
    try:
        clean_branches(dry_run=args.dry_run)
    except CmdError as e:
        print('ERROR: ' + str(e))
        sys.exit(1)