TypeTreeNode - fix from_list, impl to_dict, to_dict_list and add tests

Author: K0lb3

UnityPy/helpers/TypeTreeNode.py

@@ -35,13 +35,28 @@ def __attrs_post_init__(self):
             self._clean_name = clean_name(self.m_Name)
 
 
+TYPETREENODE_KEYS = [
+    "m_Level",
+    "m_Type",
+    "m_Name",
+    "m_ByteSize",
+    "m_Version",
+    "m_Children",
+    "m_TypeFlags",
+    "m_VariableCount",
+    "m_Index",
+    "m_MetaFlag",
+    "m_RefTypeHash",
+]
+
+
 class TypeTreeNode(TypeTreeNodeC):
     def traverse(self) -> Iterator[TypeTreeNode]:
         stack: list[TypeTreeNode] = [self]
         while stack:
             node = stack.pop()
-            stack.extend(reversed(node.m_Children))
             yield node
+            stack.extend(reversed(node.m_Children))
 
     @classmethod
     def parse(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode:
@@ -132,7 +147,7 @@ def from_list(cls, nodes: List[dict]) -> TypeTreeNode:
 
         for node in nodes:
             if isinstance(node, dict):
-                node = cls(node)
+                node = cls(**node)
 
             if node.m_Level > prev.m_Level:
                 stack.append(parent)
@@ -224,6 +239,22 @@ def dump_structure(self, indent: str = "  ") -> str:
             sb.append(child.dump_structure(indent + "  "))
         return "\n".join(sb)
 
+    def to_dict(self) -> dict:
+        return {
+            key: value
+            for key, value in ((key, getattr(self, key)) for key in TYPETREENODE_KEYS)
+            if value is not None
+        }
+
+    def to_dict_list(self) -> List[dict]:
+        return [
+            self.to_dict(),
+            *(item for child in self.m_Children for item in child.to_dict_list()),
+        ]
+
+    def __eq__(self, other: TypeTreeNode) -> bool:
+        return self.to_dict() == other.to_dict() and self.m_Children == other.m_Children
+
 
 COMMONSTRING_CACHE: Dict[Optional[UnityVersion], Dict[int, str]] = {}
 
@@ -284,3 +315,10 @@ def clean_name(name: str) -> str:
     if name[0].isdigit():
         name = f"x{name}"
     return name
+
+
+__all__ = (
+    "TypeTreeNode",
+    "get_common_strings",
+    "clean_name",
+)

tests/test_typetree.py

@@ -197,6 +197,16 @@ def test_class_node_clz():
     assert re_value == TEST_CLASS_NODE_OBJ
 
 
+def test_node_from_list_clz():
+    node = TypeTreeNode.from_list(list(TEST_CLASS_NODE.traverse()))
+    assert node == TEST_CLASS_NODE
+
+
+def test_node_from_list_dict():
+    node = TypeTreeNode.from_list(TEST_CLASS_NODE.to_dict_list())
+    assert node == TEST_CLASS_NODE
+
+
 if __name__ == "__main__":
     for x in list(locals()):
         if str(x)[:4] == "test":