Skip to content

Commit ff2d036

Browse files
K0lb3K0lb3
committed
ObjectReader - cleanup & dump_structure reimplementation
1 parent dc5a2c9 commit ff2d036

2 files changed

Lines changed: 77 additions & 34 deletions

File tree

‎UnityPy/files/ObjectReader.py‎

Lines changed: 41 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
from typing import TYPE_CHECKING, Generic, TypeVar, Union
1+
from typing import TYPE_CHECKING, Generic, TypeVar, Union, List, Optional
22
from ..enums import ClassIDType
33

44
from ..streams import EndianBinaryReader, EndianBinaryWriter
55
from ..helpers import TypeTreeHelper
66
from ..helpers.Tpk import get_typetree_nodes
7+
from ..helpers.TypeTreeNode import TypeTreeNode
78
from ..exceptions import TypeTreeError
89

910
if TYPE_CHECKING:
@@ -177,45 +178,38 @@ def __repr__(self):
177178
#
178179
###################################################
179180

180-
def dump_typetree(self, nodes: list = None) -> str:
181+
def dump_typetree_structure(
182+
self,
183+
nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
184+
indent: str = " ",
185+
) -> str:
186+
node = self._get_typetree_node(nodes)
187+
return node.dump_structure(indent=indent)
188+
189+
def read_typetree(
190+
self,
191+
nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
192+
wrap: bool = False,
193+
) -> Union[dict, T]:
181194
self.reset()
182-
sb = []
183-
nodes = self.get_typetree_nodes(nodes)
184-
TypeTreeHelper.read_typetree_str(sb, nodes, self)
185-
return "".join(sb)
186-
187-
def dump_typetree_structure(self) -> str:
188-
return TypeTreeHelper.dump_typetree(self.get_typetree_nodes())
189-
190-
def get_typetree_nodes(self, nodes: list = None) -> list:
191-
if nodes:
192-
return nodes
193-
194-
if self.serialized_type:
195-
nodes = self.serialized_type.nodes
196-
if not nodes:
197-
nodes = get_typetree_nodes(self.class_id, self.version)
198-
if not nodes:
199-
raise TypeTreeError("There are no TypeTree nodes for this object.")
200-
return nodes
201-
202-
def read_typetree(self, nodes: list = None, wrap: bool = False) -> Union[dict, T]:
203-
self.reset()
204-
nodes = self.get_typetree_nodes(nodes)
195+
node = self._get_typetree_node(nodes)
205196
ret = TypeTreeHelper.read_typetree(
206-
nodes, self.reader, as_dict=not wrap, assetsfile=self.assets_file
197+
node, self.reader, as_dict=not wrap, assetsfile=self.assets_file
207198
)
208199
if wrap:
209200
ret.set_object_reader(self)
210201
return ret
211202

212203
def save_typetree(
213-
self, tree: dict, nodes: list = None, writer: EndianBinaryWriter = None
204+
self,
205+
tree: dict,
206+
nodes: Optional[Union[TypeTreeNode, List[dict]]] = None,
207+
writer: EndianBinaryWriter = None,
214208
):
215-
nodes = self.get_typetree_nodes(nodes)
209+
node = self._get_typetree_node(nodes)
216210
if not writer:
217211
writer = EndianBinaryWriter(endian=self.reader.endian)
218-
TypeTreeHelper.write_typetree(tree, nodes, writer)
212+
TypeTreeHelper.write_typetree(tree, node, writer)
219213
data = writer.bytes
220214
self.set_raw_data(data)
221215
return data
@@ -226,3 +220,21 @@ def get_raw_data(self) -> bytes:
226220
ret = self.reader.read_bytes(self.byte_size)
227221
self.Position = pos
228222
return ret
223+
224+
def _get_typetree_node(
225+
self, nodes: Optional[Union[TypeTreeNode, List[dict]]] = None
226+
) -> TypeTreeNode:
227+
if isinstance(nodes, TypeTreeNode):
228+
return nodes
229+
elif isinstance(nodes, list):
230+
return TypeTreeNode.from_list(nodes)
231+
elif nodes is not None:
232+
raise ValueError("nodes must be a list[dict] or TypeTreeNode")
233+
234+
if self.serialized_type:
235+
nodes = self.serialized_type.nodes
236+
if not nodes:
237+
nodes = get_typetree_nodes(self.class_id, self.version)
238+
if not nodes:
239+
raise TypeTreeError("There are no TypeTree nodes for this object.")
240+
return nodes

‎UnityPy/helpers/TypeTreeNode.py‎

Lines changed: 36 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,8 @@ def traverse(self) -> Iterator[TypeTreeNode]:
3636
def parse(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode:
3737
# stack approach is way faster than recursion
3838
# using a fake root node to avoid special case for root node
39-
dummy_node = TypeTreeNode(-1, "", "", 0, 0, 0, [])
40-
dummy_root = TypeTreeNode(-1, "", "", 0, 0, 0, [dummy_node])
39+
dummy_node = cls(-1, "", "", 0, 0, 0, [])
40+
dummy_root = cls(-1, "", "", 0, 0, 0, [dummy_node])
4141

4242
stack: List[Tuple[TypeTreeNode, int]] = [(dummy_root, 1)]
4343
while stack:
@@ -47,7 +47,7 @@ def parse(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode:
4747
else:
4848
stack[-1] = (parent, count - 1)
4949

50-
node = TypeTreeNode(
50+
node = cls(
5151
m_Level=parent.m_Level + 1,
5252
m_Type=reader.read_string_to_null(),
5353
m_Name=reader.read_string_to_null(),
@@ -87,13 +87,13 @@ def read_string(reader: EndianBinaryReader, value: int) -> str:
8787
offset = value & 0x7FFFFFFF
8888
return CommonString.get(offset, str(offset))
8989

90-
fake_root: TypeTreeNode = TypeTreeNode(-1, "", "", 0, 0, 0, [])
90+
fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, 0, [])
9191
stack: List[TypeTreeNode] = [fake_root]
9292
parent = fake_root
9393
prev = fake_root
9494

9595
for raw_node in node_struct.iter_unpack(struct_data):
96-
node = TypeTreeNode(
96+
node = cls(
9797
**dict(zip(keys[:3], raw_node[:3])),
9898
**dict(zip(keys[5:], raw_node[5:])),
9999
m_Type=read_string(stringbuffer_reader, raw_node[3]),
@@ -112,6 +112,28 @@ def read_string(reader: EndianBinaryReader, value: int) -> str:
112112

113113
return fake_root.m_Children[0]
114114

115+
@classmethod
116+
def from_list(cls, nodes: List[dict]) -> TypeTreeNode:
117+
fake_root: TypeTreeNode = cls(-1, "", "", 0, 0, 0, [])
118+
stack: List[TypeTreeNode] = [fake_root]
119+
parent = fake_root
120+
prev = fake_root
121+
122+
for dict_node in nodes:
123+
node = cls(**dict_node)
124+
125+
if node.m_Level > prev.m_Level:
126+
stack.append(parent)
127+
parent = prev
128+
elif node.m_Level < prev.m_Level:
129+
while node.m_Level <= parent.m_Level:
130+
parent = stack.pop()
131+
132+
parent.m_Children.append(node)
133+
prev = node
134+
135+
return fake_root.m_Children[0]
136+
115137
def dump(self, writer: EndianBinaryWriter, version: int):
116138
stack: list[TypeTreeNode] = [self]
117139
while stack:
@@ -181,6 +203,15 @@ def write_node(node: TypeTreeNode):
181203
writer.write(node_writer.bytes)
182204
writer.write(string_writer.bytes)
183205

206+
def dump_structure(self, indent: str = " ") -> str:
207+
# dump structure similar to https://github.com/AssetRipper/TypeTreeDumps/blob/main/StructsDump
208+
sb = [
209+
f"{indent}{self.m_Type} {self.m_Name} // ByteSize{{{self.m_ByteSize:X}}}, Index{{{self.m_Index}}}, Version{{{self.m_Version}}}, TypeFlags{{{self.m_TypeFlags}}}, MetaFlag{{{self.m_MetaFlag}}}"
210+
]
211+
for child in self.m_Children:
212+
sb.append(child.dump_structure(indent + " "))
213+
return "\n".join(sb)
214+
184215

185216
COMMONSTRING_CACHE: Dict[Optional[UnityVersion], Dict[int, str]] = {}
186217

0 commit comments

Comments
 (0)