yaml_utils.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. import ruamel.yaml
  2. import sys
  3. import re
  4. import warnings
  5. # Suppress the specific warning from ruamel.yaml about floats without dots
  6. warnings.simplefilter('ignore', ruamel.yaml.error.MantissaNoDotYAML1_1Warning)
  7. def convert_to_plain_python_types(data):
  8. """
  9. Recursively converts ruamel.yaml specific types into plain Python types.
  10. Also handles converting boolean keys to strings to avoid YAML parsing issues.
  11. """
  12. if isinstance(data, ruamel.yaml.comments.CommentedMap):
  13. new_dict = {}
  14. for k, v in data.items():
  15. key = k
  16. if isinstance(k, bool):
  17. # This is a specific workaround for the 'y' key in vectors being parsed as True
  18. key = 'y' if k is True else str(k)
  19. new_dict[key] = convert_to_plain_python_types(v)
  20. return new_dict
  21. elif isinstance(data, ruamel.yaml.comments.CommentedSeq):
  22. return [convert_to_plain_python_types(i) for i in data]
  23. elif isinstance(data, ruamel.yaml.scalarstring.ScalarString):
  24. return str(data)
  25. elif isinstance(data, ruamel.yaml.anchor.Anchor):
  26. return data.value
  27. return data
  28. def load_unity_yaml(file_path):
  29. """
  30. Safely loads a Unity YAML file, which may contain multiple documents,
  31. into a list of Python objects.
  32. """
  33. try:
  34. with open(file_path, 'r', encoding='utf-8') as f:
  35. yaml_content = f.read()
  36. # Use regex to fix the Unity-specific shorthand tags.
  37. # This looks for "--- !u!### &..." and replaces it with a valid tag format.
  38. processed_content = re.sub(r'--- !u!(\d+) &', r'--- !<tag:unity3d.com,2011:\1> &', yaml_content)
  39. # Unity files start with a %YAML directive, which ruamel.yaml can handle.
  40. # The documents are separated by '---'.
  41. yaml = ruamel.yaml.YAML(typ='rt') # Use round-trip mode to handle custom tags
  42. # allow_duplicate_keys=True is important for some Unity files
  43. yaml.allow_duplicate_keys = True
  44. documents = list(yaml.load_all(processed_content))
  45. return documents
  46. except FileNotFoundError:
  47. print(f"Error: File not found at {file_path}", file=sys.stderr)
  48. return []
  49. except ruamel.yaml.YAMLError as e:
  50. print(f"Error parsing YAML file {file_path}: {e}", file=sys.stderr)
  51. return []
  52. except Exception as e:
  53. print(f"An unexpected error occurred while reading {file_path}: {e}", file=sys.stderr)
  54. return []
  55. if __name__ == '__main__':
  56. # Example usage:
  57. # This allows testing the module directly.
  58. # You would run this from the root of the project as:
  59. # python Assets/LLM/source/utils/yaml_utils.py <path_to_some_asset>.asset
  60. if len(sys.argv) > 1:
  61. file_to_load = sys.argv[1]
  62. print(f"Attempting to load Unity YAML file: {file_to_load}")
  63. docs = load_unity_yaml(file_to_load)
  64. if docs:
  65. print(f"Successfully loaded {len(docs)} documents.")
  66. for i, doc in enumerate(docs):
  67. print(f"--- Document {i+1} ---")
  68. # A simple print might not show all nested details, but it's a good check.
  69. print(doc)
  70. print()
  71. else:
  72. print("Failed to load documents or the file was empty.")
  73. else:
  74. print("Please provide a file path as a command-line argument to test the loader.")