prateek · February 28, 2020 20:16
diff --git a/golang_import_order_cleaner.py b/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()
	#!/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()