Generates a Go vendor directory in a bazel workspace

genvendor.py

#!/usr/bin/env python3

import ast
import collections
import os
import subprocess
import sys
import time


Repo = collections.namedtuple('Repo', ['name', 'build_file', 'commit', 'remote'])


class WorkspaceASTFileVisitor(ast.NodeVisitor):
    def __init__(self):
        self.repos = []
    
    def visit_Call(self, node):
        if node.func.id == 'new_git_repository':
            args = {}
            for kw in node.keywords:
                args[kw.arg] = kw.value.s
            self.repos.append(Repo(**args))


class BuildfileASTFileVisitor(ast.NodeVisitor):
    def __init__(self):
        self._basename = ''
        self._prefix = ''
        self._name = ''
    
    @property
    def basename(self):
        if self._basename != '':
            return self._basename
        return self._prefix + '/' + self._name
    
    def visit_Call(self, node):
        if node.func.id.startswith('external_go_package'):
            for kw in node.keywords:
                if kw.arg == 'base_pkg':
                    self._basename = kw.value.s
        if node.func.id == 'go_prefix':
            if len(node.args) == 1:
                self._prefix = node.args[0].s
        if node.func.id in ['cgo_library', 'go_library']:
            for kw in node.keywords:
                if kw.arg == 'name':
                    self._name = kw.value.s


def get_repo_basename(workspace_dir, repo):
    buildfile_path = os.path.join(workspace_dir, repo.build_file)
    if not os.path.exists(buildfile_path):
        exit("Buildfile " + buildfile_path + " is missing.")
    buildfile = open(buildfile_path).read()
    node = ast.parse(buildfile)
    visitor = BuildfileASTFileVisitor()
    visitor.visit(node)
    return visitor.basename

def get_repos(workspace_dir):
    workspace_file_path = os.path.join(workspace_dir, 'WORKSPACE')
    if not os.path.exists(workspace_file_path):
        exit("Workspace is expected to contain a WORKSPACE file")
    workspace = open(workspace_file_path).read()
    node = ast.parse(workspace)
    visitor = WorkspaceASTFileVisitor()
    visitor.visit(node)
    return visitor.repos

def stream_output(*cmd):
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    while p.poll() is None:
        stdout = p.stdout.readline()
        if stdout:
            sys.stdout.writelines([stdout.decode('utf8')])
        stderr = p.stderr.readline()
        if stderr:
            sys.stderr.writelines([stderr.decode('utf8')])
    stdout, stderr = p.communicate()
    if stdout:
        sys.stdout.write(stdout.decode('utf8'))
    if stderr:
        sys.stderr.write(stderr.decode('utf8'))
    
    if p.returncode > 0:
        exit("Error executing command")

def gen_vendor(workspace_dir):
    repos = get_repos(workspace_dir)
    for repo in repos:
        basename = get_repo_basename(workspace_dir, repo)
        repo_dir = os.path.join(workspace_dir, 'vendor', basename)
        if not os.path.exists(repo_dir):
            print("Cloning " + repo.name)
            os.makedirs(repo_dir)
            stream_output('git', 'clone', repo.remote, '--recursive', repo_dir)
        os.chdir(repo_dir)
        print("Updating " + repo.name)
        stream_output('git', 'fetch', '--all')
        stream_output('git', 'checkout', repo.commit)


if __name__ == '__main__':
    if len(sys.argv) != 2:
        exit("Usage: genvendor.py <path to workspace>")
    gen_vendor(os.path.abspath(os.path.expanduser(sys.argv[1])))