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

import sys
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
from file_utils import create_guid_to_path_map

def parse_scene_or_prefab(file_path, guid_map=None, assets_dir=None):
    """
    Parses a Unity scene or prefab file and returns a flat list of
    GameObject dictionaries with embedded component data.
    """
    if guid_map is None or assets_dir is None:
        p = Path(file_path).resolve()
        while p.name != 'Assets' and p.parent != p:
            p = p.parent
        if p.name == 'Assets':
            assets_dir = str(p)
        else:
            assets_dir = str(Path(file_path).resolve().parent.parent)
        guid_map = create_guid_to_path_map(assets_dir)

    documents = load_unity_yaml(file_path)
    if not documents:
        return None

    print("    -> Converting YAML documents to Python objects...")
    raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}
    object_map = {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}
    print(f"    -> Found {len(object_map)} objects in the file.")

    flat_gameobject_list = []
    transform_to_gameobject = {}
    parent_to_children_transforms = {}

    # First pass: Populate relationship maps
    print("    -> Pass 1/2: Mapping object relationships...")
    for file_id, obj_data in object_map.items():
        if 'Transform' in obj_data:
            parent_id = obj_data['Transform'].get('m_Father', {}).get('fileID')
            if parent_id:
                if parent_id not in parent_to_children_transforms:
                    parent_to_children_transforms[parent_id] = []
                parent_to_children_transforms[parent_id].append(file_id)
    
    for file_id, obj_data in object_map.items():
        if 'GameObject' in obj_data:
            for comp in obj_data['GameObject'].get('m_Component', []):
                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: Process each GameObject
    print("    -> Pass 2/2: Extracting GameObject data and components...")
    total_gos = sum(1 for obj in object_map.values() if 'GameObject' in obj)
    processed_gos = 0
    for file_id, obj_data in object_map.items():
        if 'GameObject' in obj_data:
            processed_gos += 1
            if processed_gos % 100 == 0:
                print(f"        ...processed {processed_gos}/{total_gos} GameObjects...")

            go_info = obj_data['GameObject']
            
            source_obj_info = go_info.get('m_CorrespondingSourceObject')
            if source_obj_info and source_obj_info.get('guid'):
                guid = source_obj_info['guid']
                source_prefab_path = guid_map.get(guid)
                if source_prefab_path:
                    modifications = go_info.get('m_Modification')
                    flat_gameobject_list.append({
                        'fileID': file_id,
                        'm_Name': go_info.get('m_Name'),
                        'isPrefab': True,
                        'sourcePrefabGUID': guid,
                        'm_Modification': modifications
                    })
                    continue

            components_data = []
            for comp_ref in go_info.get('m_Component', []):
                comp_id = comp_ref['component']['fileID']
                if comp_id in object_map:
                    component_data = object_map[comp_id]
                    # Create a new structure to avoid modifying the original object_map
                    restructured_component = {}
                    for component_type, component_values in component_data.items():
                        if isinstance(component_values, dict):
                            # Copy the inner dict and inject the fileID
                            new_values = component_values.copy()
                            new_values['fileID'] = comp_id
                            restructured_component[component_type] = new_values
                        else:
                            # Fallback for non-dict components, though unlikely
                            restructured_component[component_type] = component_values
                    components_data.append(restructured_component)
            
            children_ids = []
            my_transform_id = None
            for t_id, go_id in transform_to_gameobject.items():
                if go_id == file_id:
                    my_transform_id = t_id
                    break
            
            if my_transform_id in parent_to_children_transforms:
                child_transform_ids = parent_to_children_transforms[my_transform_id]
                for child_t_id in child_transform_ids:
                    child_go_id = transform_to_gameobject.get(child_t_id)
                    if child_go_id:
                        children_ids.append(child_go_id)

            flat_gameobject_list.append({
                '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_ComponentsData': components_data,
                'm_Children': children_ids
            })

    return flat_gameobject_list