Rearranges the source files in a Unity package folder for asmdef creation

asmdefify.py

#!/usr/bin/env python3

"""
Rearranges the source files in a Unity project folder for asmdef creation.
This script will create two directories within the target directory, "Runtime" and "Editor".
All files that live within an existing "Editor" folder will be moved to this new top-level
"Editor" folder, but with the rest of their folder hierarchy maintained. All other files
will be moved into the new top-level "Runtime" folder, also with their existing folder hierarchy
maintained.

For example, this folder structure:

Game/
- GameClass1.cs
- GameClass2.cs
- Editor/
-- EditorClass1.cs
- ANamespace/
-- GameClass3.cs
-- Editor/
--- EditorClass2.cs

will be transformed to:

Game/
- Runtime/
-- GameClass1.cs
-- GameClass2.cs
-- ANamespace/
--- GameClass3.cs
- Editor/
-- EditorClass1.cs
--- ANamespace/
---- EditorClass2.cs
"""

import argparse
import logging
import os
import shutil
import sys
from typing import NamedTuple, List, Iterable

logging.basicConfig(level=logging.INFO)
LOG = logging.getLogger('')

EDITOR = 'Editor'
RUNTIME = 'Runtime'


class PathMapping(NamedTuple):
    cur_path: str
    new_path: str


def is_editor_dir(dir_path: str) -> bool:
    return EDITOR in dir_path.split(os.sep)


def is_subpath(path: str, superpaths: Iterable[str]) -> bool:
    """True of the given path is a subpath of any of the
    potential superpaths"""
    for sp in superpaths:
        if os.path.commonprefix([path, sp]) == sp:
            return True
    return False


def is_source_file(file_path) -> bool:
    return os.path.splitext(file_path)[1] == '.cs'


def get_editor_dirs(root_dir: str) -> Iterable[str]:
    """Yields all Editor directories in the given root_dir"""
    for (dir_path, dirs, files) in os.walk(root_dir, topdown=True):
        dir_path_rel = os.path.relpath(dir_path, root_dir)
        if is_editor_dir(dir_path_rel):
            # Don't need to into children
            dirs.clear()
            yield dir_path


def has_source_files(dir_path: str) -> bool:
    """True if the given directory has any source files in it"""
    for filename in os.listdir(dir_path):
        if os.path.isfile(filename) and is_source_file(filename):
            return True
    return False


def build_path_mappings(root_dir: str) -> List[PathMapping]:
    LOG.info(f'Processing {root_dir}...')
    editor_dirs = list(get_editor_dirs(root_dir))
    root_editor_dir = None
    non_root_editor_dirs = []
    for editor_dir in editor_dirs:
        path_rel = os.path.relpath(editor_dir, root_dir)
        if path_rel == EDITOR:
            root_editor_dir = editor_dir
        else:
            non_root_editor_dirs.append(editor_dir)

    # We can early-exit:
    # - If we have no editor directories, or
    # - If we have a single root editor dir, and its parent directory
    #   _doesn't_ have any source files in it
    if len(non_root_editor_dirs) == 0:
        if root_editor_dir is None:
            return []
        root_editor_parent = os.path.dirname(root_editor_dir)
        if not has_source_files(root_editor_parent):
            return []

    # Build our PathMappings
    mappings: List[PathMapping] = []
    for (dir_path, dirs, files) in os.walk(root_dir, topdown=True):
        # Determine the directory to move the files into
        rel_dir_path = os.path.relpath(dir_path, root_dir)
        if is_subpath(dir_path, editor_dirs):
            # create a new path; move the 'Editor' path component to the front
            components = rel_dir_path.split(os.sep)
            components.remove(EDITOR)
            components.insert(0, EDITOR)
            new_dir_path = os.path.join(root_dir, os.sep.join(components))
        else:
            components = rel_dir_path.split(os.sep)
            if components[0] != RUNTIME:
                components.insert(0, RUNTIME)
            new_dir_path = os.path.join(root_dir, os.sep.join(components))

        if dir_path == new_dir_path:
            continue

        for file in files:
            cur_path = os.path.join(dir_path, file)
            new_path = os.path.join(new_dir_path, file)

            if file == f'{EDITOR}.meta':
                if dir_path == root_dir:
                    # Don't move top-level Editor.meta
                    continue
                else:
                    # Delete all other Editor.metas
                    new_path = None
            elif file == f'{RUNTIME}.meta' and dir_path == root_dir:
                # Don't move top-level Runtime.meta
                continue

            mappings.append(PathMapping(cur_path, new_path))

    return mappings


def apply_path_mappings(root_dir: str, mappings: List[PathMapping]) -> None:
    LOG.info(f'Moving {len(mappings)} files...')
    for mapping in mappings:
        if mapping.new_path is None:
            LOG.info(f'DELETING {os.path.relpath(mapping.cur_path, root_dir)}')
            os.remove(mapping.cur_path)
        else:
            new_dir = os.path.split(mapping.new_path)[0]
            os.makedirs(new_dir, exist_ok=True)
            LOG.info(f'{os.path.relpath(mapping.cur_path, root_dir)} -> '
                     f'{os.path.relpath(mapping.new_path, root_dir)}')
            shutil.move(mapping.cur_path, mapping.new_path)

    LOG.info(f'Cleaning empty directories...')
    for (dir_path, dirs, files) in os.walk(root_dir, topdown=False):
        if len(os.listdir(dir_path)) == 0:
            LOG.info(f'Removing {os.path.relpath(dir_path, root_dir)}')
            os.rmdir(dir_path)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('directory', help='Directory to organize')
    ap.add_argument('-n', '--dry-run', action='store_true',
                    help='Don\'t move files, just print results')
    args = ap.parse_args()

    root_dir = os.path.abspath(args.directory)
    if not os.path.isdir(root_dir):
        LOG.error(f'No such directory: {root_dir}')
        sys.exit(1)

    mappings = build_path_mappings(root_dir)
    if len(mappings) == 0:
        LOG.info('Nothing to do!')
        sys.exit(0)

    if args.dry_run:
        mapping_strings = ['']  # An empty string for `join`
        for mapping in mappings:
            src = os.path.relpath(mapping.cur_path, root_dir)
            if mapping.new_path is None:
                dst = 'DELETE'
            else:
                dst = os.path.relpath(mapping.new_path, root_dir)
            mapping_strings.append(f'{src} -> {dst}')
        joined = '\n  '.join(mapping_strings)
        LOG.info(f'Would move {len(mappings)} files:{joined}')
    else:
        apply_path_mappings(root_dir, mappings)


if __name__ == '__main__':
    main()