Unity3D CLI utilities

unitool.py

#!/usr/bin/env python3

import argparse
import os
import re
import shutil
import subprocess
import sys
from ruamel.yaml import YAML
from ruamel.yaml.parser import Parser

Parser.DEFAULT_TAGS['!u!'] = 'tag:unity3d.com,2011:'

LIBRARY_CACHE=os.path.expanduser('~/.unitool/Libraries')
DOCUMENTS = {}

yaml = YAML(typ='rt')

def _reconstruct_document(loader, suffix, node):
    document = loader.construct_mapping(node, True)
    print(loader.anchors)
    return document

#yaml.constructor.BaseConstructor.add_multi_constructor('tag:unity3d.com,2011:', _reconstruct_document)

################ Platform switching handling ####################
def _rsync_directories(src, dst):
    command = f'rsync -avz "{src}" "{dst}"'
    p = subprocess.Popen(command, shell=True)
    exitCode = p.wait()
    return exitCode == 0


def switch_platform_with_rsync(args):
    args = vars(args)
    
    target_platform = args['to-platform']
    project_path = os.path.abspath(args['unity-project'])
    path_to_project_cache = os.path.join(LIBRARY_CACHE, project_path.replace('/', '_'))
    path_to_backup = os.path.join(path_to_project_cache, args['from-platform'])
    path_to_restore = os.path.join(path_to_project_cache, target_platform, 'Library')
    path_to_library = os.path.join(project_path, 'Library')

    if not os.path.isdir(path_to_library) or len(os.listdir(path_to_library)) == 0:
        print('Library is empty, skipping...')
    else:
        if not os.path.isdir(path_to_backup):
            print(f'Did not find any cache, performing full copy "{path_to_library}" to "{path_to_backup}"...')
        else:
            print(f'Syncing current backup with Library ("{path_to_library}" to "{path_to_backup}")...')
        
        if not _rsync_directories(path_to_library, path_to_backup):
            print(f'Backup failed! Aborting!')
            sys.exit(100)
        
        print('Done!')

    if os.path.isdir(path_to_library):
        print(f'Remove {path_to_library}')
        shutil.rmtree(path_to_library)

    if not os.path.isdir(path_to_restore):
        print('No cache found, creating empty Library...')
        os.makedirs(path_to_library)
    else:
        print(f'Restoring Library from {path_to_restore} to {path_to_library}')
        if not _rsync_directories(path_to_restore, project_path):
            print(f'Restoring failed! Aborting!')
            sys.exit(101)
    
    print(f'!!! Remember to launch Unity Editor via Unity Hub and select platform {target_platform}!!!')


def switch_platform_with_mv(args):
    args = vars(args)
    
    target_platform = args['to-platform']
    project_path = os.path.abspath(args['unity-project'])
    path_to_project_cache = os.path.join(LIBRARY_CACHE, project_path.replace('/', '_'))
    path_to_backup = os.path.join(path_to_project_cache, args['from-platform'], 'Library')
    path_to_restore = os.path.join(path_to_project_cache, target_platform, 'Library')
    path_to_library = os.path.join(project_path, 'Library')
    
    if not os.path.isdir(path_to_library) or len(os.listdir(path_to_library)) == 0:
        print('Library is empty, skipping...')
    else:
        if os.path.isdir(path_to_backup):
            shutil.rmtree(path_to_backup)
        
        print(f'Moving Library to backup location ({path_to_library} to {path_to_backup})')
        subprocess.run(["mv", path_to_library, path_to_backup])
    
    if not os.path.isdir(path_to_restore):
        print('Nothing in cache, creating empty Library')
        os.makedirs(path_to_library)
    else:
        print(f'Moving cache Library to proper location ({path_to_restore} to {path_to_library})')
        subprocess.run(["mv", path_to_restore, path_to_library])


def backup_library(args):
    args = vars(args)
    
    target_platform = args['platform']
    project_path = os.path.abspath(args['unity-project'])
    path_to_project_cache = os.path.join(LIBRARY_CACHE, project_path.replace('/', '_'))
    path_to_backup = os.path.join(path_to_project_cache, target_platform)
    path_to_library = os.path.join(project_path, 'Library')

    os.makedirs(path_to_backup)
    print(f'Syncing current backup with Library ("{path_to_library}" to "{path_to_backup}")...')
            
    if not _rsync_directories(path_to_library, path_to_backup):
        print(f'Backup failed! Aborting!')
        sys.exit(102)
    
    print('Done!')


def checkout_platform(args):
    args = vars(args)
    
    target_platform = args['platform']
    project_path = os.path.abspath(args['unity-project'])
    path_to_project_cache = os.path.join(LIBRARY_CACHE, project_path.replace('/', '_'))
    path_to_backup = os.path.join(path_to_project_cache, target_platform)
    path_to_library = os.path.join(project_path, 'Library')

    os.makedirs(path_to_library)
    print(f'Restoring backup ("{path_to_backup}" to "{path_to_library}")...')
            
    if not _rsync_directories(path_to_backup, project_path):
        print(f'Checkout failed! Aborting!')
        sys.exit(103)
    
    print('Done!')


################ Files related operations ###################
def _assert_unity_project_directory():
    directory_content = os.listdir('.')

    if 'Library' not in directory_content or 'Assets' not in directory_content or 'ProjectSettings' not in directory_content:
        print('Run this command from top-level project directory!')
        sys.exit(200)
 
def inspect_file(args):
    _assert_unity_project_directory()
    args = vars(args)
    
    lookup_name = args['file']
    lookup_extension = os.path.splitext(lookup_name)[1]
    find_result = subprocess.run(['find', '.', '-iname', lookup_name], capture_output=True, check=True, encoding='utf-8')
    paths = list(filter(None, find_result.stdout.split('\n')))
    if not paths:
        print('File not found')
        sys.exit(201)

    if len(paths) > 1:
        print('Found multiple matching files. Which one:')
        for i, path in enumerate(paths):
            print(f'{i+1}) {path}')
        
        idx_string = input('Type index, nothing to abort: ')
        try:
            idx = int(idx_string)
        except ValueError:
            sys.exit(0)
        
        file_path = paths[idx - 1]
    else:
        file_path = paths[0]
    
    with open(f'{file_path}.meta', 'r') as f:
        meta_content = yaml.load(f)
        guid = meta_content['guid']

        print(f'Inspecting {file_path} ({guid})...')
    
    occurrences_result = subprocess.run(['grep', '-r', guid, '.'], capture_output=True, check=True, encoding='utf-8')
    paths = list(filter(None, occurrences_result.stdout.split('\n')))
    occurrences = []

    for path in paths:
        path = path.split(':')[0]
        if path == f'{file_path}.meta' or not path.startswith('./'):
            continue
        occurrences.append(path)
        
    if not occurrences:
        print('Found no Unity references.')
    else:
        for occurrence in occurrences:
            print(f'Checking {occurrence} for references...')
            file_name = os.path.basename(occurrence)
            name, extension = os.path.splitext(file_name)

            if extension == ".prefab" or extension == ".unity":
                with open(occurrence, 'r') as f:
                    text = f.read()
                    text = re.sub(r' stripped$', '', text, flags=re.MULTILINE)
                    content = list(yaml.load_all(text))
                
                if lookup_extension == ".cs":
                    id_to_document = {}
                    for document in content:
                        id_to_document[int(document.anchor.value)] = document
                    
                    for document in content:
                        if "MonoBehaviour" not in document:
                            continue
                        
                        if document['MonoBehaviour']['m_Script']['guid'] != guid:
                            continue
                        
                        if "m_GameObject" in document['MonoBehaviour']:
                            game_object = id_to_document[document['MonoBehaviour']['m_GameObject']['fileID']]
                            game_object_name = game_object['GameObject']['m_Name']
                            print(f'Attached to GameObject: "{game_object_name}"')
                        elif "m_PrefabInternal" in document['MonoBehaviour']:
                            prefab = id_to_document[document['MonoBehaviour']['m_PrefabInternal']['fileID']]
                            prefab_guid = prefab['Prefab']['m_SourcePrefab']['guid']
                            prefix = prefab_guid[0:2]
                            with open(os.path.join('Library', 'metadata', prefix, f'{prefab_guid}.info'), 'r', errors='replace') as f:
                                lines = f.read().splitlines()[1:-1]
                                lines.insert(0, 'Info:')
                                prefab_info = yaml.load("\n".join(lines))
                                prefab_name = prefab_info['Info']['mainRepresentation']['name']
                            print(f'Part of Prefab: "{prefab_name}"')



######################## CLI handling ########################

parser = argparse.ArgumentParser(prog='unitool', description='Various tools to help working with Unity3D')

subparsers = parser.add_subparsers(help=None, dest='command')

parser_switch_platform = subparsers.add_parser('switch-platform', help='Switch Library to pre-imported one (if exists)')
parser_switch_platform.add_argument('unity-project', type=str, help='Path to Unity project')
parser_switch_platform.add_argument('from-platform', type=str, help='Name of platform which should backed up')
parser_switch_platform.add_argument('to-platform', type=str, help='Name of platform which should become active')
parser_switch_platform.set_defaults(func=switch_platform_with_mv)

parser_backup_library = subparsers.add_parser('backup-library', help='Update backup of current Library for faster switch')
parser_backup_library.add_argument('unity-project', type=str, help='Path to Unity project')
parser_backup_library.add_argument('platform', type=str, help='Name of platform to backup')
parser_backup_library.set_defaults(func=backup_library)

parser_backup_library = subparsers.add_parser('checkout-platform', help='Restore Library content from cache without syncing current platform\'s backup')
parser_backup_library.add_argument('unity-project', type=str, help='Path to Unity project')
parser_backup_library.add_argument('platform', type=str, help='Name of platform to check out')
parser_backup_library.set_defaults(func=checkout_platform)

parser_inspect = subparsers.add_parser('inspect', help='Find all available information about given file')
parser_inspect.add_argument('file', type=str, help='File name to inspect, including extension.')
parser_inspect.set_defaults(func=inspect_file)

args = parser.parse_args()

if not args.command:
    parser.print_help()
    sys.exit(1)
else:
    args.func(args)