import ordering cleanup

golang_import_order_cleaner.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# golang_import_order_cleaner.py: cleans up import orders in golang files per the m3db convention
#
# Warning: This is very much alpha, take a backup of your files (e.g. commit to git) before running this.
#
# example usage:
# (1) Update all *.go files in the local directory (except vendor/), in-place
# $ golang_import_order_cleaner.py -o inplace
#
# (2) Update all the files in a particular git branch
# find and update all the *.go files modified in the git branch (assuming master is base)
# $ git diff master..HEAD --name-only | grep -e '.go$' | golang_import_order_cleaner.py -i - -o inplace

import fileinput
import argparse
import fnmatch
import os
import sys
from os import path

def get_local_packages():
    for (dirpath, dirnames, filenames) in os.walk('.'):
        return [d for d in dirnames if d != 'vendor' and d[0] != '.']

def get_go_files(dirs):
    matches = set()
    for d in dirs:
        for root, dirnames, filenames in os.walk(d):
            for filename in fnmatch.filter(filenames, '*.go'):
                matches.add(os.path.join(root, filename))
    return list(matches)

def cleanup_imports_and_return(imports):
    os_packages = []
    thirdparty_packages = []
    for i in imports:
        if i.strip() == "":
            continue
        elif i.find(".com") != -1 or i.find(".net") != -1 or i.find(".org") != -1 or i.find(".in") != -1:
            thirdparty_packages.append(i)
        else:
            os_packages.append(i)

    l = ["import ("]
    needs_new_line = False
    if os_packages:
        l.extend(os_packages)
        needs_new_line = True
    if thirdparty_packages:
        if needs_new_line:
            l.append("")
        l.extend(thirdparty_packages)
    l.append(")")
    return l

def parse_go_file(f):
    with open(f, 'r') as go_file:
        lines = [i.rstrip() for i in go_file.readlines()]
        in_import_block = False
        imports = []
        output_lines = []
        for line in lines:
            if in_import_block:
                endIdx = line.find(")")
                if endIdx != -1:
                    in_import_block = False
                    output_lines.extend(cleanup_imports_and_return(imports))
                    imports = []
                    continue
                imports.append(line)
            else:
                importIdx = line.find("import (")
                if importIdx != -1:
                    in_import_block = True
                    continue
                output_lines.append(line)
        output_lines.append("")
        return "\n".join(output_lines)


def main():
    parser = argparse.ArgumentParser(
        description='Tool to make cleaning up import orders easily')

    parser.add_argument('-o', '--output', default='stdout',
                        choices=['inplace', 'stdout'],
                        help='output target [default: stdout]')

    parser.add_argument('-i', '--input',
                        help='file with list of go src files to operate upon, operates on all non-vendor dirs by default')

    args = parser.parse_args()
    output = args.output

    go_files = []
    input = args.input
    if not input:
        print >>sys.stderr, "No input specified, operating upon all *.go files in local dir"
        dirs = get_local_packages()
        go_files = get_go_files(dirs)
    else:
      go_files_list = fileinput.input(files=(input))
      go_files = [i.strip() for i in go_files_list]

    for f in go_files:
        # print >>sys.stderr, "Parsing ", f
        parsed = parse_go_file(f)
        if output == "stdout":
            print parsed
        else:
            with open(f, 'w') as ofile:
                ofile.write(parsed)
            print >>sys.stderr, "Updated: %s" % f

if __name__ == '__main__':
    main()