gerrit cherry-pick script

#!/usr/bin/env python
# repo cherry-pick script

import sys, os
import subprocess
from xml.etree import ElementTree
from StringIO import StringIO
import json

GIT = 'git'
SSH = 'ssh'

REPO = 'repo'
REPO_START = REPO, 'start'
REPO_ABANDON = REPO, 'abandon'

def get_top():
    return os.getenv('TOP') or os.getenv('ANDROID_BUILD_TOP')

class Manifest:
    def __init__(self):
        self.pathmap = {}
        self.remotes = {}
        self.default = ''
        self.load()

        if self.default in self.remotes:
            self.git_server = self.remotes[self.default].get('fetch')
            self.review_server = self.remotes[self.default].get('review')
        else:
            print 'no default remote server'
            sys.exit(-1)

        for scheme in ('git://', 'http://', 'https://'):
            if self.git_server.startswith(scheme):
                print 'unsupported gerrit transport type.'
                print 'this script supports only ssh transport type'
                sys.exit(-1)

        if self.git_server.startswith('ssh://'):
            if self.git_server.endswith('/'):
                self.git_server = self.git_server[:-1]
        elif self.git_server == '..':
            # aosp does
            print 'unsupported gerrit transport type.'
            print 'this script supports only ssh transport type'
            sys.exit(-1)
        elif not self.git_server.endswith(':'):
            self.git_server += ':'

        # XXX
        self.review_server = self.git_server
        if self.review_server.startswith('ssh://'):
            self.review_server = self.review_server[6:]

    def load(self):
        for filename in ('manifest.xml', 'local_manifest.xml'):
            manifest = os.path.sep.join(('.repo', filename))
            if not os.access(manifest, os.R_OK):
                continue
            p = ElementTree.parse(open(manifest))
            r = p.getroot()

            for remote in r.findall('remote'):
                self.remotes[remote.get('name')] = remote

            for default in r.findall('default'):
                self.default = default.get('remote')

            for project in r.findall('remove-project'):
                name = project.get('name')
                if name in self.pathmap:
                    del self.pathmap[name]

            for project in r.findall('project'):
                self.pathmap[project.get('name')] = project.get('path')

    def __getitem__(self, k):
        return self.pathmap.get(k)

    def __contains__(self, k):
        return k in self.pathmap

class CherryPick:
    def __init__(self, branch = None):
        self.manifest = Manifest()
        self.git_server = self.manifest.git_server
        self.review_server = self.manifest.review_server
        self.applied = []
        self.branch = branch

    def reset_branch(self, project, directory):
        print 'reset project %s with branch %s'%(project, self.branch)
        cmd = REPO_ABANDON + (self.branch, directory)
        subprocess.call(cmd)
        cmd = REPO_START + (self.branch, directory)
        if subprocess.call(cmd) != 0:
            print 'branch create failed. abort'
            sys.exit(-1)
        self.applied.append(directory)

    def apply(self, project, ref):
        if not project in self.manifest:
            print 'no such project %s. abort'%project
            sys.exit(-1)

        directory = os.path.sep.join((get_top(), self.manifest[project]))

        if self.branch and not directory in self.applied:
            self.reset_branch(project, directory)

        print 'fetch item %s'%ref
        cmd = GIT, 'fetch', '/'.join((self.git_server, project)), ref
        if subprocess.call(cmd, cwd = directory):
            print 'fetch failed'
            sys.exit(-1)

        print 'apply item %s'%ref
        cmd = GIT, 'cherry-pick', 'FETCH_HEAD'
        if subprocess.call(cmd, cwd = directory):
            print 'apply failed'
            sys.exit(-1)

    def query(self, items):
        query = ('change:' + x for x in items)
        gerrit_cmd = ('gerrit', 'query', ' OR '.join(query), '--current-patch-set', '--format', 'json')
        cmd = (SSH, self.review_server, ' '.join(gerrit_cmd))

        print 'query items to server'
        result = StringIO(subprocess.Popen(cmd, stdout=subprocess.PIPE, close_fds=True).stdout.read())
        r = {}
        for line in result:
            item = json.loads(line)
            if not 'status' in item:
                continue

            number = item['number']

            if not item['open']:
                print 'item %s is not opened. abort'%number
                sys.exit(-1)

            r[number] = item
        return r

    def main(self, items):
        result = self.query(items)
        for number in items:
            item = result[number]
            project = item['project']
            ref = item['currentPatchSet']['ref']
            rev = item['currentPatchSet']['revision']
            self.apply(project, ref)

if __name__ == '__main__':
    if len(sys.argv) == 1:
        print 'Usage: cherry-pick item1 item2 item3 ... [branch]'
        sys.exit(-1)

    top = get_top()
    if not top:
        print 'TOP or ANDROID_BUILD_TOP is not defined. abort'
        sys.exit(-1)

    os.chdir(top)

    if sys.argv[-1].isdigit():
        items = sys.argv[1:]
        branch = None
    else:
        items = sys.argv[1:-1]
        branch = sys.argv[-1]

    cherryPick = CherryPick(branch)
    cherryPick.main(items)