fix: use fixed TypeTreeGeneratorAPI and correctly use it to generate nodes, also be more explicit regarding related errors

Author: K0lb3

UnityPy/files/ObjectReader.py

@@ -282,10 +282,7 @@ def _get_typetree_node(
         if not node:
             node = get_typetree_node(self.class_id, self.version)
             if node.m_Type == "MonoBehaviour":
-                try:
-                    node = self._try_monobehaviour_node(node)
-                except Exception:
-                    pass
+                node = self.generate_monobehaviour_node(node)
         if not node:
             raise TypeTreeError("There are no TypeTree nodes for this object.")
         return node
@@ -297,21 +294,29 @@ def parse_as_object(self, node: Optional[NodeInput] = None, check_read: bool = T
     def parse_as_dict(self, node: Optional[NodeInput] = None, check_read: bool = True) -> dict[str, Any]:
         return self.read_typetree(nodes=node, wrap=False, check_read=check_read)  # type: ignore
 
-    def _try_monobehaviour_node(self, base_node: TypeTreeNode) -> TypeTreeNode:
+    def parse_monobehaviour_head(self, mb_node: Optional[TypeTreeNode] = None) -> MonoBehaviour:
+        if mb_node is None:
+            mb_node = get_typetree_node(ClassIDType.MonoBehaviour, self.version)
+
+        mb = self.read_typetree(nodes=mb_node, wrap=True, check_read=False)
+        return cast(MonoBehaviour, mb)
+
+    def generate_monobehaviour_node(self, mb_node: Optional[TypeTreeNode] = None) -> TypeTreeNode:
         env = self.assets_file.environment
         generator = env.typetree_generator
         if generator is None:
-            raise ValueError("No typetree generator set!")
-        monobehaviour = cast(MonoBehaviour, self.parse_as_object(base_node, check_read=False))
+            raise ValueError("MonoBehaviour detected, but no typetree_generator set to the environment!")
+
+        monobehaviour = self.parse_monobehaviour_head(mb_node)
         script = monobehaviour.m_Script.deref_parse_as_object()
 
         if script.m_Namespace != "":
             fullname = f"{script.m_Namespace}.{script.m_ClassName}"
         else:
             fullname = script.m_ClassName
 
-        node = generator.get_nodes_up(base_node, script.m_AssemblyName, fullname)
+        node = generator.get_nodes_up(script.m_AssemblyName, fullname)
         if node:
             return node
         else:
-            raise ValueError("Failed to get custom MonoBehaviour node!")
+            raise ValueError(f"Failed to generate MonoBehaviour node for {fullname} of {script.m_AssemblyName}!")

UnityPy/helpers/TypeTreeGenerator.py

@@ -4,7 +4,7 @@
 from .TypeTreeNode import TypeTreeNode
 
 try:
-    from TypeTreeGeneratorAPI import TypeTreeGenerator as TypeTreeGeneratorBase
+    from TypeTreeGeneratorAPI import TypeTreeGenerator as TypeTreeGeneratorBase  # pyright: ignore[reportAssignmentType]
 except ImportError:
 
     class TypeTreeGeneratorBase:
@@ -45,50 +45,32 @@ def load_local_dll_folder(self, dll_dir: str):
                 data = f.read()
                 self.load_dll(data)
 
-    def get_nodes_up(self, base_node: TypeTreeNode, assembly: str, fullname: str) -> TypeTreeNode:
-        root = self.cache.get((assembly, fullname))
-        if root is not None:
-            return root
+    def get_nodes_up(self, assembly: str, fullname: str) -> TypeTreeNode:
+        key = (assembly, fullname)
+        if key in self.cache:
+            return self.cache[key]
 
         if not assembly.endswith(".dll"):
             assembly = f"{assembly}.dll"
+
         base_nodes = self.get_nodes(assembly, fullname)
 
-        base_root = base_nodes[0]
-        root = TypeTreeNode(
-            base_root.m_Level,
-            base_root.m_Type,
-            base_root.m_Name,
-            0,
-            0,
-            m_MetaFlag=base_root.m_MetaFlag,
-            m_Children=base_node.m_Children[:],
+        node = TypeTreeNode.from_list(
+            [
+                TypeTreeNode(
+                    base_node.m_Level,
+                    base_node.m_Type,
+                    base_node.m_Name,
+                    0,
+                    0,
+                    m_MetaFlag=base_node.m_MetaFlag,
+                )
+                for base_node in base_nodes
+            ]
         )
-        stack: List[TypeTreeNode] = []
-        parent = root
-        prev = root
-
-        for base_node in base_nodes[1:]:
-            node = TypeTreeNode(
-                base_node.m_Level,
-                base_node.m_Type,
-                base_node.m_Name,
-                0,
-                0,
-                m_MetaFlag=base_node.m_MetaFlag,
-            )
-            if node.m_Level > prev.m_Level:
-                stack.append(parent)
-                parent = prev
-            elif node.m_Level < prev.m_Level:
-                while node.m_Level <= parent.m_Level:
-                    parent = stack.pop()
-
-            parent.m_Children.append(node)
-            prev = node
-
-        self.cache[(assembly, fullname)] = root
-        return root
+
+        self.cache[key] = node
+        return node
 
 
 __all__ = ("TypeTreeGenerator",)

pyproject.toml

@@ -59,7 +59,7 @@ UnityPy = "UnityPy.cli:main"
 
 [project.optional-dependencies]
 # optional dependencies must be lowercase/normalized
-ttgen = ["typetreegeneratorapi>=0.0.7"]
+ttgen = ["typetreegeneratorapi>=0.0.10"]
 full = ["unitypy[ttgen]"]
 tests = ["pytest", "pillow", "psutil", "unitypy[full]"]
 dev = ["ruff", "unitypy[tests]"]