checkout-pr.py

import sys
import subprocess
import re
import shlex
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('command')
parser.add_argument('action', choices=['pull', 'push'])
args = parser.parse_args()

branch = args.command.split()[-1]
remote = args.command.split()[2]
user = re.search(r'git@github.com:(.*)/pandas.git', args.command).group(1)

if args.action == 'pull':
    subprocess.run(['git', 'checkout', '-b', f'pr/{user}/{branch}'])
    subprocess.run(['git', 'checkout', f'pr/{user}/{branch}'])
    subprocess.run(['git', 'pull', remote, branch])
    subprocess.run(['git', 'reset', '--hard', 'FETCH_HEAD'])
else:
    subprocess.run(['git', 'push', remote, f'HEAD:{branch}'])