christiandeange · June 21, 2018 21:25
diff --git a/fix_sources.py b/fix_sources.py
 #!/usr/bin/env python
 #
 # Script to fix module roots that are incorrectly labelled as source roots by
 # `buck project`. It iterates over each .iml file in the project and modifies
 # source folder entries to ensure that the module root itself is not
 # registered as a source root. There are several issues this script resolves:
 #  - Module root is marked as source root
 #  - `src/test` is marked as test source root (instead of `src/test/java`)
 #  - Folders holding protobuf definitions are marked as excluded
 #
 # The `xmltodict` pip module is required for this to run.
 # This script can be executed using either python2.7 or python3.
 #

 from __future__ import print_function

 import argparse
 import fnmatch
 import glob
 import logging
 import os
 import sys
 import time

 try:
    import xmltodict
 except ImportError:
    print('Please install the xmltodict module: `pip install xmltodict`')
    exit(1)

 NO_SOURCE = 'file://$MODULE_DIR$'

 logging.basicConfig(format='%(message)s', level=logging.DEBUG)


 class Result(object):
    """Model class to hold result of parsing a module source."""
    OK = 0
    FIXED = 1
    DELETE = 2
    UNKNOWN = 3


 class ImlFile(object):
    """Holds a reference to a the backing .iml file for a module in the project."""
    def __init__(self, project_root, module_root):
        self.project_root = project_root
        self.module_root = module_root
        self.modifications = False

    def repofile(self, folder):
        """Returns a string representing the relative path to a folder from the project root."""
        if folder.startswith(self.project_root):
            folder = folder[len(self.project_root):]
        return folder

    def url_to_folder(self, url):
        """Returns the reference to a folder in .iml-notation, relative to the module's root."""
        relative_to_module = url[len(NO_SOURCE) + 1:]
        return os.path.realpath(os.path.join(self.module_root, relative_to_module))

    def has_file_with_extension(self, folder, extension):
        """
        Recursively checks if any file with the given extension exists in the folder,
        or any of the folder's subfolders.
        """
        for _, _, files in os.walk(folder):
            if any(filter(lambda f: f.endswith(extension), files)):
                return True
        return False

    def remap_source_folder(self, folders, folder_type, source_folder, new_relative_path):
        """
        Performs the appropriate action to remap this source root entry to the proper folder.
        Used to correctly label both normal source roots and test source roots.
        """
        new_module_folder = os.path.join(self.module_root, new_relative_path)
        new_module_source = os.path.join(NO_SOURCE, new_relative_path)

        if any(filter(lambda node: node['@url'] == new_module_source, folders)):
            # "fixed" entry already exists in iml: just delete this one
            logging.debug('%s should be unmarked as source root' % self.repofile(self.module_root))
            return Result.DELETE
        elif os.path.isdir(new_module_folder):
            # found an incorrectly-labelled source root, point it to the proper folder
            logging.debug('{module_root} should have {folder_type} source root {new_folder}'.format(
                module_root=self.repofile(self.module_root),
                folder_type=folder_type,
                new_folder=self.repofile(new_module_folder)
            ))

            source_folder['@url'] = new_module_source
            source_folder.pop('@packagePrefix', None)
            return Result.FIXED
        else:
            # found a source root definition that cannot be resolved
            logging.debug('Unresolved module: %s' % self.repofile(self.module_root))
            return Result.UNKNOWN

    def process_source_folder(self, folders, source_folder):
        """Checks a single source folder definition for this module for correctness."""
        url = source_folder['@url']

        if url == NO_SOURCE:
            # fix instances where the module root is marked as source root
            folder_type = 'test' if test_source(source_folder) else 'main'
            new_relative_path = os.path.join('src', folder_type, 'java')
            return self.remap_source_folder(folders, folder_type, source_folder, new_relative_path)

        if url == os.path.join(NO_SOURCE, 'src', 'test'):
            # fix instances where `src/test` is marked as test root (instead of `src/test/java`)
            new_relative_path = os.path.join('src', 'test', 'java')
            return self.remap_source_folder(folders, 'test', source_folder, new_relative_path)

        # this source folder is already valid
        return Result.OK

    def process_exclude_folder(self, folders, exclude_folder):
        """
        Deletes excluded folder definitions if they contain protobuf definitions.
        This allows protobuf defintions to show up in file searches.
        """
        url = exclude_folder['@url']
        folder_path = self.url_to_folder(url)

        if self.has_file_with_extension(folder_path, '.proto'):
            logging.debug('%s should be unexcluded' % self.repofile(folder_path))
            return Result.DELETE
        return Result.OK

    def process_folders(self, content, folder_type, folder_handler):
        """
        Inspect folder definition for folder_type nodes in the module.
        Each folder definition is then passed to the folder_handler function,
        which returns a Result class to indicate the new status of said definition.
        """
        original_folders = as_list(content.get(folder_type))
        new_folders = []
        for folder in original_folders:
            result = folder_handler(original_folders, folder)
            self.modifications |= (result == Result.FIXED or result == Result.DELETE)
            if result != Result.DELETE:
                new_folders.append(folder)

        if len(new_folders) == 1:
            content[folder_type] = new_folders[0]
        else:
            content[folder_type] = new_folders

    def process_content(self, content):
        """Inspect source folder and excluded folder definitions for this module."""
        self.process_folders(content, 'sourceFolder', self.process_source_folder)
        self.process_folders(content, 'excludeFolder', self.process_exclude_folder)

    def process_component(self, component):
        for content in as_list(component.get('content')):
            if content['@url'] == NO_SOURCE:
                self.process_content(content)

    def process_module(self, module):
        for component in as_list(module.get('component')):
            if component['@name'] == 'NewModuleRootManager':
                self.process_component(component)

    def process(self, iml_contents):
        for module in as_list(iml_contents.get('module')):
            if module['@type'] == 'JAVA_MODULE':
                self.process_module(module)
        return iml_contents


 def as_list(obj):
    if obj is None:
        return []
    elif type(obj) is list:
        return obj
    else:
        return [obj]


 def test_source(node):
    return node['@isTestSource'] == 'true'


 def process_xml(args, iml_file_path):
    """
    Inspect a single .iml file and resolve any incorrect folder definitions.
    If this .iml file had any definitions that needed to be corrected, it will
    be modified on the filesystem, and will return `true`.
    """
    project_root = os.path.realpath(args.project_dir)
    module_root = os.path.dirname(iml_file_path)
    iml = ImlFile(project_root, module_root)

    # read iml contents
    with open(iml_file_path, 'r') as fd:
        iml_contents = xmltodict.parse(fd.read())

    # modify it to resolve incorrect links
    iml_contents = iml.process(iml_contents)

    # write changes back if necessary
    if iml.modifications and not args.dryrun:
        with open(iml_file_path, 'w') as fd:
            xmltodict.unparse(iml_contents, output=fd, pretty=True)

    return iml.modifications

 def process_folder(args, folder, in_root_folder=False):
    """
    Recursively process a folder on the filesystem to check for modules.
    If there are modules found in this folder, process them to check for
    folder definition correctness. Yields booleans to indicate if a
    module .iml was modified to correct its folder definitions.
    """
    imls = glob.glob('%s/*.iml' % folder)

    # modules cannot be nested within another module, so once an .iml file is
    # found in a folder, the folder's subdirectories don't need to be walked
    # (the only exception to this is the project folder)
    if imls and not in_root_folder:
        for iml in imls:
            # yields the number of modifications made
            yield process_xml(args, iml)
    else:
        for filename in os.listdir(folder):
            relative_path = os.path.join(folder, filename)
            if in_root_folder and filename == 'buck-out':
                # skip ./buck-out, it's a huge folder with no imls
                continue
            if os.path.isdir(relative_path):
                yield sum(process_folder(args, relative_path))


 def main():
    parser = argparse.ArgumentParser(description='Fix incorrect folder definitions in .iml files caused by `buck project`.')
    parser.add_argument('-p', '--project-dir', dest='project_dir', action='store', default='.', help='The root folder of the repo')
    parser.add_argument('-n', '--dry-run', dest='dryrun', action='store_true', help='Do not modify anything; just report what would change')
    parser.add_argument('-v', '--verbose', dest='verbose', action='store_true', help='Print diagnostic info for .iml processing')
    args = parser.parse_args()

    if args.dryrun:
        print('Running in dry run mode, no modifications will be made.')

    project_folder = args.project_dir
    if not os.path.isdir(os.path.join(project_folder, '.idea')):
        print('No .idea folder found in %s' % project_folder)
        exit(1)

    if not args.verbose:
        logging.disable(logging.DEBUG)

    print('Fixing source folders:\r', end='')
    sys.stdout.flush()

    start_time = time.time()
    modifications = sum(process_folder(args, project_folder, in_root_folder=True))
    end_time = time.time()

    print('Fixing %d source folders: finished in %.1f sec' % (modifications, end_time - start_time))

 if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass
	#!/usr/bin/env python
	#
	# Script to fix module roots that are incorrectly labelled as source roots by
	# `buck project`. It iterates over each .iml file in the project and modifies
	# source folder entries to ensure that the module root itself is not
	# registered as a source root. There are several issues this script resolves:
	# - Module root is marked as source root
	# - `src/test` is marked as test source root (instead of `src/test/java`)
	# - Folders holding protobuf definitions are marked as excluded
	#
	# The `xmltodict` pip module is required for this to run.
	# This script can be executed using either python2.7 or python3.
	#

	from __future__ import print_function

	import argparse
	import fnmatch
	import glob
	import logging
	import os
	import sys
	import time

	try:
	import xmltodict
	except ImportError:
	print('Please install the xmltodict module: `pip install xmltodict`')
	exit(1)

	NO_SOURCE = 'file://$MODULE_DIR$'

	logging.basicConfig(format='%(message)s', level=logging.DEBUG)


	class Result(object):
	"""Model class to hold result of parsing a module source."""
	OK = 0
	FIXED = 1
	DELETE = 2
	UNKNOWN = 3


	class ImlFile(object):
	"""Holds a reference to a the backing .iml file for a module in the project."""
	def __init__(self, project_root, module_root):
	self.project_root = project_root
	self.module_root = module_root
	self.modifications = False

	def repofile(self, folder):
	"""Returns a string representing the relative path to a folder from the project root."""
	if folder.startswith(self.project_root):
	folder = folder[len(self.project_root):]
	return folder

	def url_to_folder(self, url):
	"""Returns the reference to a folder in .iml-notation, relative to the module's root."""
	relative_to_module = url[len(NO_SOURCE) + 1:]
	return os.path.realpath(os.path.join(self.module_root, relative_to_module))

	def has_file_with_extension(self, folder, extension):
	"""
	Recursively checks if any file with the given extension exists in the folder,
	or any of the folder's subfolders.
	"""
	for _, _, files in os.walk(folder):
	if any(filter(lambda f: f.endswith(extension), files)):
	return True
	return False

	def remap_source_folder(self, folders, folder_type, source_folder, new_relative_path):
	"""
	Performs the appropriate action to remap this source root entry to the proper folder.
	Used to correctly label both normal source roots and test source roots.
	"""
	new_module_folder = os.path.join(self.module_root, new_relative_path)
	new_module_source = os.path.join(NO_SOURCE, new_relative_path)

	if any(filter(lambda node: node['@url'] == new_module_source, folders)):
	# "fixed" entry already exists in iml: just delete this one
	logging.debug('%s should be unmarked as source root' % self.repofile(self.module_root))
	return Result.DELETE
	elif os.path.isdir(new_module_folder):
	# found an incorrectly-labelled source root, point it to the proper folder
	logging.debug('{module_root} should have {folder_type} source root {new_folder}'.format(
	module_root=self.repofile(self.module_root),
	folder_type=folder_type,
	new_folder=self.repofile(new_module_folder)
	))

	source_folder['@url'] = new_module_source
	source_folder.pop('@packagePrefix', None)
	return Result.FIXED
	else:
	# found a source root definition that cannot be resolved
	logging.debug('Unresolved module: %s' % self.repofile(self.module_root))
	return Result.UNKNOWN

	def process_source_folder(self, folders, source_folder):
	"""Checks a single source folder definition for this module for correctness."""
	url = source_folder['@url']

	if url == NO_SOURCE:
	# fix instances where the module root is marked as source root
	folder_type = 'test' if test_source(source_folder) else 'main'
	new_relative_path = os.path.join('src', folder_type, 'java')
	return self.remap_source_folder(folders, folder_type, source_folder, new_relative_path)

	if url == os.path.join(NO_SOURCE, 'src', 'test'):
	# fix instances where `src/test` is marked as test root (instead of `src/test/java`)
	new_relative_path = os.path.join('src', 'test', 'java')
	return self.remap_source_folder(folders, 'test', source_folder, new_relative_path)

	# this source folder is already valid
	return Result.OK

	def process_exclude_folder(self, folders, exclude_folder):
	"""
	Deletes excluded folder definitions if they contain protobuf definitions.
	This allows protobuf defintions to show up in file searches.
	"""
	url = exclude_folder['@url']
	folder_path = self.url_to_folder(url)

	if self.has_file_with_extension(folder_path, '.proto'):
	logging.debug('%s should be unexcluded' % self.repofile(folder_path))
	return Result.DELETE
	return Result.OK

	def process_folders(self, content, folder_type, folder_handler):
	"""
	Inspect folder definition for folder_type nodes in the module.
	Each folder definition is then passed to the folder_handler function,
	which returns a Result class to indicate the new status of said definition.
	"""
	original_folders = as_list(content.get(folder_type))
	new_folders = []
	for folder in original_folders:
	result = folder_handler(original_folders, folder)
	self.modifications \|= (result == Result.FIXED or result == Result.DELETE)
	if result != Result.DELETE:
	new_folders.append(folder)

	if len(new_folders) == 1:
	content[folder_type] = new_folders[0]
	else:
	content[folder_type] = new_folders

	def process_content(self, content):
	"""Inspect source folder and excluded folder definitions for this module."""
	self.process_folders(content, 'sourceFolder', self.process_source_folder)
	self.process_folders(content, 'excludeFolder', self.process_exclude_folder)

	def process_component(self, component):
	for content in as_list(component.get('content')):
	if content['@url'] == NO_SOURCE:
	self.process_content(content)

	def process_module(self, module):
	for component in as_list(module.get('component')):
	if component['@name'] == 'NewModuleRootManager':
	self.process_component(component)

	def process(self, iml_contents):
	for module in as_list(iml_contents.get('module')):
	if module['@type'] == 'JAVA_MODULE':
	self.process_module(module)
	return iml_contents


	def as_list(obj):
	if obj is None:
	return []
	elif type(obj) is list:
	return obj
	else:
	return [obj]


	def test_source(node):
	return node['@isTestSource'] == 'true'


	def process_xml(args, iml_file_path):
	"""
	Inspect a single .iml file and resolve any incorrect folder definitions.
	If this .iml file had any definitions that needed to be corrected, it will
	be modified on the filesystem, and will return `true`.
	"""
	project_root = os.path.realpath(args.project_dir)
	module_root = os.path.dirname(iml_file_path)
	iml = ImlFile(project_root, module_root)

	# read iml contents
	with open(iml_file_path, 'r') as fd:
	iml_contents = xmltodict.parse(fd.read())

	# modify it to resolve incorrect links
	iml_contents = iml.process(iml_contents)

	# write changes back if necessary
	if iml.modifications and not args.dryrun:
	with open(iml_file_path, 'w') as fd:
	xmltodict.unparse(iml_contents, output=fd, pretty=True)

	return iml.modifications

	def process_folder(args, folder, in_root_folder=False):
	"""
	Recursively process a folder on the filesystem to check for modules.
	If there are modules found in this folder, process them to check for
	folder definition correctness. Yields booleans to indicate if a
	module .iml was modified to correct its folder definitions.
	"""
	imls = glob.glob('%s/*.iml' % folder)

	# modules cannot be nested within another module, so once an .iml file is
	# found in a folder, the folder's subdirectories don't need to be walked
	# (the only exception to this is the project folder)
	if imls and not in_root_folder:
	for iml in imls:
	# yields the number of modifications made
	yield process_xml(args, iml)
	else:
	for filename in os.listdir(folder):
	relative_path = os.path.join(folder, filename)
	if in_root_folder and filename == 'buck-out':
	# skip ./buck-out, it's a huge folder with no imls
	continue
	if os.path.isdir(relative_path):
	yield sum(process_folder(args, relative_path))


	def main():
	parser = argparse.ArgumentParser(description='Fix incorrect folder definitions in .iml files caused by `buck project`.')
	parser.add_argument('-p', '--project-dir', dest='project_dir', action='store', default='.', help='The root folder of the repo')
	parser.add_argument('-n', '--dry-run', dest='dryrun', action='store_true', help='Do not modify anything; just report what would change')
	parser.add_argument('-v', '--verbose', dest='verbose', action='store_true', help='Print diagnostic info for .iml processing')
	args = parser.parse_args()

	if args.dryrun:
	print('Running in dry run mode, no modifications will be made.')

	project_folder = args.project_dir
	if not os.path.isdir(os.path.join(project_folder, '.idea')):
	print('No .idea folder found in %s' % project_folder)
	exit(1)

	if not args.verbose:
	logging.disable(logging.DEBUG)

	print('Fixing source folders:\r', end='')
	sys.stdout.flush()

	start_time = time.time()
	modifications = sum(process_folder(args, project_folder, in_root_folder=True))
	end_time = time.time()

	print('Fixing %d source folders: finished in %.1f sec' % (modifications, end_time - start_time))

	if __name__ == '__main__':
	try:
	main()
	except KeyboardInterrupt:
	pass