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

import ruamel.yaml
import sys
import re
import warnings

# Suppress the specific warning from ruamel.yaml about floats without dots
warnings.simplefilter('ignore', ruamel.yaml.error.MantissaNoDotYAML1_1Warning)

def convert_to_plain_python_types(data):
    """
    Recursively converts ruamel.yaml specific types into plain Python types.
    Also handles converting boolean keys to strings to avoid YAML parsing issues.
    """
    if isinstance(data, ruamel.yaml.comments.CommentedMap):
        new_dict = {}
        for k, v in data.items():
            key = k
            if isinstance(k, bool):
                # This is a specific workaround for the 'y' key in vectors being parsed as True
                key = 'y' if k is True else str(k)
            new_dict[key] = convert_to_plain_python_types(v)
        return new_dict
    elif isinstance(data, ruamel.yaml.comments.CommentedSeq):
        return [convert_to_plain_python_types(i) for i in data]
    elif isinstance(data, ruamel.yaml.scalarstring.ScalarString):
        return str(data)
    elif isinstance(data, ruamel.yaml.anchor.Anchor):
        return data.value
    return data

def load_unity_yaml(file_path):
    """
    Safely loads a Unity YAML file, which may contain multiple documents,
    into a list of Python objects.
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            yaml_content = f.read()

        # Use regex to fix the Unity-specific shorthand tags.
        # This looks for "--- !u!### &..." and replaces it with a valid tag format.
        processed_content = re.sub(r'--- !u!(\d+) &', r'--- !<tag:unity3d.com,2011:\1> &', yaml_content)
        # Remove non-standard keywords like 'stripped' from document headers
        processed_content = re.sub(r'(&\d+)\s+stripped', r'\1', processed_content)
        
        # Unity files start with a %YAML directive, which ruamel.yaml can handle.
        # The documents are separated by '---'.
        yaml = ruamel.yaml.YAML(typ='rt') # Use round-trip mode to handle custom tags
        # allow_duplicate_keys=True is important for some Unity files
        yaml.allow_duplicate_keys = True
        yaml.strict = False
        
        documents = list(yaml.load_all(processed_content))
        return documents
    except FileNotFoundError:
        print(f"Error: File not found at {file_path}", file=sys.stderr)
        return []
    except ruamel.yaml.YAMLError as e:
        print(f"Error parsing YAML file {file_path}: {e}", file=sys.stderr)
        return []
    except Exception as e:
        print(f"An unexpected error occurred while reading {file_path}: {e}", file=sys.stderr)
        return []

if __name__ == '__main__':
    # Example usage:
    # This allows testing the module directly.
    # You would run this from the root of the project as:
    # python Assets/LLM/source/utils/yaml_utils.py <path_to_some_asset>.asset
    if len(sys.argv) > 1:
        file_to_load = sys.argv[1]
        print(f"Attempting to load Unity YAML file: {file_to_load}")
        
        docs = load_unity_yaml(file_to_load)
        
        if docs:
            print(f"Successfully loaded {len(docs)} documents.")
            for i, doc in enumerate(docs):
                print(f"--- Document {i+1} ---")
                # A simple print might not show all nested details, but it's a good check.
                print(doc)
                print()
        else:
            print("Failed to load documents or the file was empty.")
    else:
        print("Please provide a file path as a command-line argument to test the loader.")