terra-llm/Assets/LLM/source/utils/virtual_tree_builder.py

import sys
import json
import io
from pathlib import Path
import ruamel.yaml

# Add the utils directory to the Python path
utils_path = Path(__file__).parent
sys.path.append(str(utils_path))

from yaml_utils import load_unity_yaml, convert_to_plain_python_types

def build_virtual_tree(file_path):
    """
    Parses a Unity scene or prefab file and builds a nested dictionary 
    representing the GameObject hierarchy.
    """
    documents = load_unity_yaml(file_path)
    if not documents:
        return None

    # First, build the map from the raw documents to preserve anchors
    raw_object_map = {doc.anchor.value: doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}

    # Now, convert the mapped objects to plain Python types
    object_map = {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}

    nodes = {}
    transform_to_gameobject = {}

    # First pass: create nodes for all GameObjects and map transforms
    for file_id, obj_data in object_map.items():
        if 'GameObject' in obj_data:
            go_info = obj_data['GameObject']
            nodes[file_id] = {
                'fileID': file_id,
                'm_Name': go_info.get('m_Name'),
                'm_IsActive': go_info.get('m_IsActive', 1),
                'm_TagString': go_info.get('m_TagString'),
                'm_Layer': go_info.get('m_Layer'),
                'm_Components': [comp['component']['fileID'] for comp in go_info.get('m_Component', []) if comp and 'component' in comp],
                'm_Children': []
            }
            # Find the transform component to map it back to the GameObject
            for comp in go_info.get('m_Component', []):
                if comp and 'component' in comp:
                    comp_id = comp['component']['fileID']
                    if comp_id in object_map and 'Transform' in object_map.get(comp_id, {}):
                        transform_to_gameobject[comp_id] = file_id
                        break

    # Second pass: build the hierarchy
    root_nodes = list(nodes.values())
    for go_id, node in nodes.items():
        # Find the transform for the current GameObject
        transform_id = None
        for comp_id in node['m_Components']:
            if comp_id in transform_to_gameobject:
                transform_id = comp_id
                break
        
        if transform_id:
            transform_data = object_map.get(transform_id, {}).get('Transform', {})
            if transform_data:
                parent_transform_id = transform_data.get('m_Father', {}).get('fileID')

                if parent_transform_id and parent_transform_id != 0:
                    parent_go_id = transform_to_gameobject.get(parent_transform_id)
                    if parent_go_id and parent_go_id in nodes:
                        # Check to prevent adding a node to its own children list
                        if nodes[parent_go_id] is not node:
                            nodes[parent_go_id]['m_Children'].append(node)
                            if node in root_nodes:
                                root_nodes.remove(node)

    return root_nodes