git-cleanup

#!/usr/bin/env python3

import subprocess
import sys


def run_cmd(cmd):
    return subprocess.check_output(cmd).decode('utf8').strip()


def get_current_branch():
    return run_cmd(['git', 'name-rev', '--name-only', 'HEAD'])


def delete_local_branch(name):
    return run_cmd(['git', 'branch', '-d', name])


def delete_remote_branch(name):
    remote, name = name.split('/', 1)
    return run_cmd(['git', 'push', remote, ':{}'.format(name)])


def list_branches(remote=False):
    """
    Return a list of branch names that should be safe to delete. The given opts
    will be passed into `git branch` and should determine whether to list local
    or remote branches.

    Local branch output:

        $ git branch --merged
          INFRA-658-user-cleanup
        * master
          master-new
          new-master

    Remote branch output:

        $ git branch -r --merged
          origin/DISTR-609-scaffold-components-and-fetch-artifacts
          origin/HEAD -> origin/master
          origin/INFRA-658-user-cleanup
          origin/VIDPUB-136-API-GET-project-by-ID-precommit
          origin/feeder-worker
          origin/master
    """
    cmd = ['git', 'branch']
    if remote:
        cmd.append('-r')
    cmd.extend(['--merged'])

    lines = run_cmd(cmd).splitlines()
    names = [parse_branch_name(line) for line in lines]
    return [name for name in names if safe_to_delete(name, remote)]


def parse_branch_name(line):
    return line.lstrip(' *').split()[0]


def safe_to_delete(name, is_remote):
    remote_name = None
    if is_remote:
        remote_name, branch = name.split('/', 1)
    else:
        branch = name
    is_protected = branch in ('master', 'develop', 'HEAD')
    is_mine = remote_name == 'origin' or not is_remote
    return is_mine and not is_protected


def pluralize(xs, suffix='s'):
    return suffix if len(xs) > 0 else ''


def get_answer(base_prompt):
    answer_map = {'y': True, 'yes': True, 'n': False, 'no': False, 'q': False}
    prompt = f'{base_prompt} [Y/n]: '
    while True:
        answer = input(prompt).strip().lower() or 'y'
        if answer in answer_map:
            return answer_map[answer]


def main():
    current_branch = get_current_branch()
    if current_branch != 'master':
        print(f'Error: git-cleanup must be run from the `master` branch (currently: `{current_branch}`)') # noqa
        return 1

    local_branches = list_branches(remote=False)
    remote_branches = list_branches(remote=True)

    work = [
        ('local', local_branches, delete_local_branch),
        ('remote', remote_branches, delete_remote_branch),
    ]

    for kind, branches, delete in work:
        if not branches:
            continue

        count = len(branches)
        branch_word = 'branch' + pluralize(branches, 'es')
        print(f'Found {count} fully merged {kind} {branch_word}:')
        for branch in branches:
            print(f' - {branch}')
        do_delete = get_answer(f'Delete {count} {kind} {branch_word}?')
        if do_delete:
            for branch in branches:
                delete(branch)
        print()


if __name__ == '__main__':
    sys.exit(main())