Skip to content

Commit 42e1320

Browse files
K0lb3K0lb3
committed
TypeTree - use UnknownObject, move clean_name to TypeTreeNode, cache cleaned type name
1 parent 8bdaaa0 commit 42e1320

7 files changed

Lines changed: 406 additions & 252 deletions

File tree

‎UnityPy/classes/UnknownObject.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
from typing import Optional
2+
3+
from ..helpers.TypeTreeNode import TypeTreeNode
4+
from .Object import Object
5+
6+
7+
class UnknownObject(Object):
8+
"""An object of unknown type that showed up during typetree parsing."""
9+
10+
__node__: Optional[TypeTreeNode]
11+
12+
def __init__(self, __node__: TypeTreeNode = None, **kwargs):
13+
self.__node__ = __node__
14+
self.__dict__.update(**kwargs)
15+
16+
def get_type(self):
17+
return self.__node__.m_Type if self.__node__ else None
18+
19+
def __repr__(self) -> str:
20+
def format_value(v):
21+
vstr = repr(v)
22+
if len(vstr) > 100:
23+
return vstr[:97] + "..."
24+
return vstr
25+
26+
return f"<UnknownObject<{self.get_type()}> {', '.join(
27+
f'{k}={format_value(v)}' for k, v in self.__dict__.items() if k != '__node__'
28+
)}>"

‎UnityPy/classes/__init__.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
from .Object import Object as Object
2-
from .PPtr import PPtr as PPtr
31
from .generated import *
42
from .legacy_patch import *
3+
from .Object import Object as Object
4+
from .PPtr import PPtr as PPtr
5+
from .UnknownObject import UnknownObject as UnknownObject

‎UnityPy/helpers/TypeTreeHelper.py‎

Lines changed: 77 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from .TypeTreeNode import TypeTreeNode
1212

1313
Object = classes.Object
14+
UnknownObject = classes.UnknownObject
1415
PPtr = classes.PPtr
1516

1617
if TYPE_CHECKING:
@@ -134,7 +135,7 @@ def read_typetree(
134135
if expected_read and read_typetree_boost:
135136
data = reader.read_bytes(expected_read)
136137
return read_typetree_boost(
137-
data, root_node, reader.endian, as_dict, assetsfile, classes, clean_name
138+
data, root_node, reader.endian, as_dict, assetsfile, classes
138139
)
139140

140141
pos = reader.Position
@@ -218,7 +219,9 @@ def read_value(
218219
else:
219220
config = config.copy()
220221
config.has_registry = True
221-
value[child.m_Name] = read_value(child, reader, config)
222+
value[child.m_Name if config.as_dict else child._clean_name] = read_value(
223+
child, reader, config
224+
)
222225

223226
if not config.as_dict:
224227
if node.m_Type.startswith("PPtr<"):
@@ -228,21 +231,29 @@ def read_value(
228231
m_PathID=value["m_PathID"],
229232
)
230233
else:
231-
clz = getattr(classes, node.m_Type, Object)
232-
clz_kwargs = {clean_name(key): value for key, value in value.items()}
234+
clz = getattr(classes, node.m_Type, UnknownObject)
233235
try:
234-
value = clz(**clz_kwargs)
236+
value = clz(**value)
235237
except TypeError:
236-
extra_keys = set(clz_kwargs.keys()) - set(clz.__annotations__)
237-
value = clz(
238-
**{
239-
key: value
240-
for key, value in clz_kwargs.items()
241-
if key in clz.__annotations__
242-
}
243-
)
244-
for key in extra_keys:
245-
setattr(value, key, clz_kwargs[key])
238+
keys = set(value.keys())
239+
annotation_keys = set(clz.__annotations__)
240+
missing_keys = annotation_keys - keys
241+
if clz is UnknownObject or missing_keys:
242+
value = UnknownObject(node, **value)
243+
else:
244+
extra_keys = keys - annotation_keys
245+
if extra_keys:
246+
value = clz(
247+
**{
248+
key: value
249+
for key, value in value.items()
250+
if key in clz.__annotations__
251+
}
252+
)
253+
for key in extra_keys:
254+
setattr(value, key, value[key])
255+
else:
256+
value = UnknownObject(**value)
246257

247258
if align:
248259
reader.align_stream()
@@ -313,53 +324,66 @@ def read_value_array(
313324
}
314325
for _ in range(size)
315326
]
327+
elif node.m_Type.startswith("PPtr<"):
328+
value = [
329+
PPtr[Any](
330+
assetsfile=config.assetsfile,
331+
**{
332+
child.m_Name: read_value(child, reader, config)
333+
for child in node.m_Children
334+
},
335+
)
336+
for _ in range(size)
337+
]
316338
else:
317-
if node.m_Type.startswith("PPtr<"):
339+
clz = getattr(
340+
classes,
341+
node.m_Type,
342+
UnknownObject,
343+
)
344+
keys = set(child._clean_name for child in node.m_Children)
345+
annotation_keys = set(clz.__annotations__)
346+
missing_keys = annotation_keys - keys
347+
extra_keys = keys - annotation_keys
348+
if missing_keys or clz is UnknownObject:
318349
value = [
319-
PPtr[Any](
320-
assetsfile=config.assetsfile,
350+
UnknownObject(
351+
node,
321352
**{
322-
child.m_Name: read_value(child, reader, config)
353+
child._clean_name: read_value(child, reader, config)
323354
for child in node.m_Children
324355
},
325356
)
326357
for _ in range(size)
327358
]
359+
elif extra_keys:
360+
value = [None] * size
361+
for i in range(size):
362+
value_i_d = {
363+
child._clean_name: read_value(child, reader, config)
364+
for child in node.m_Children
365+
}
366+
value_i = clz(
367+
**{
368+
key: value
369+
for key, value in value_i_d.items()
370+
if key in annotation_keys
371+
}
372+
)
373+
for key in extra_keys:
374+
setattr(value_i, key, value_i_d[key])
375+
value[i] = value_i
328376
else:
329-
clz = getattr(
330-
classes,
331-
node.m_Type,
332-
Object,
333-
)
334-
clean_names = [clean_name(child.m_Name) for child in node.m_Children]
335-
if all(name in clz.__annotations__ for name in clean_names):
336-
value = [
337-
clz(
338-
**{
339-
name: read_value(child, reader, config)
340-
for name, child in zip(clean_names, node.m_Children)
341-
}
342-
)
343-
for _ in range(size)
344-
]
345-
else:
346-
extra_keys = set(clean_names) - set(clz.__annotations__)
347-
value = [None] * size
348-
for i in range(size):
349-
value_i_d = {
350-
clean_name(child.m_Name): read_value(child, reader, config)
377+
value = [
378+
clz(
379+
**{
380+
child._clean_name: read_value(child, reader, config)
351381
for child in node.m_Children
352382
}
353-
value_i = clz(
354-
**{
355-
key: value
356-
for key, value in value_i_d.items()
357-
if key in clz.__annotations__
358-
}
359-
)
360-
for key in extra_keys:
361-
setattr(value_i, key, value_i_d[key])
362-
value[i] = value_i
383+
)
384+
for _ in range(size)
385+
]
386+
363387
if align:
364388
reader.align_stream()
365389
return value
@@ -409,7 +433,7 @@ def clean_name(name: str) -> str:
409433

410434

411435
def write_value(
412-
value: Any,
436+
value: Union[dict[str, Any], Object],
413437
node: TypeTreeNode,
414438
writer: EndianBinaryWriter,
415439
config: TypeTreeConfig,
@@ -456,9 +480,7 @@ def write_value(
456480
else:
457481
config = config.copy()
458482
config.has_registry = True
459-
write_value(
460-
getattr(value, clean_name(child.m_Name)), child, writer, config
461-
)
483+
write_value(getattr(value, child._clean_name), child, writer, config)
462484

463485
if align:
464486
writer.align_stream()

‎UnityPy/helpers/TypeTreeNode.py‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
from __future__ import annotations
2+
3+
import re
24
from struct import Struct
3-
from typing import Optional, List, Tuple, TYPE_CHECKING, Dict, Iterator
5+
from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple
46

57
from attrs import define, field
68

@@ -27,6 +29,10 @@ class TypeTreeNodeC:
2729
m_Index: Optional[int] = None
2830
m_MetaFlag: Optional[int] = None
2931
m_RefTypeHash: Optional[int] = None
32+
_clean_name: str = field(init=False)
33+
34+
def __attrs_post_init__(self):
35+
self._clean_name = clean_name(self.m_Name)
3036

3137

3238
class TypeTreeNode(TypeTreeNodeC):
@@ -41,8 +47,8 @@ def traverse(self) -> Iterator[TypeTreeNode]:
4147
def parse(cls, reader: EndianBinaryReader, version: int) -> TypeTreeNode:
4248
# stack approach is way faster than recursion
4349
# using a fake root node to avoid special case for root node
44-
dummy_node = cls(-1, "", "", 0, 0, 0, [])
45-
dummy_root = cls(-1, "", "", 0, 0, 0, [dummy_node])
50+
dummy_node = cls(-1, "", "", 0, 0, [])
51+
dummy_root = cls(-1, "", "", 0, 0, [dummy_node])
4652

4753
stack: List[Tuple[TypeTreeNode, int]] = [(dummy_root, 1)]
4854
while stack:
@@ -262,3 +268,17 @@ def _get_blob_node_struct(endian: str, version: int) -> tuple[Struct, list[str]]
262268
keys.append("m_RefTypeHash")
263269

264270
return Struct(struct_type), keys
271+
272+
273+
def clean_name(name: str) -> str:
274+
# keep in sync with TypeTreeHelper.cpp
275+
if name.startswith("(int&)"):
276+
name = name[6:]
277+
if name.endswith("?"):
278+
name = name[:-1]
279+
name = re.sub(r"[ \.:\-\[\]]", "_", name)
280+
if name in ["pass", "from"]:
281+
name += "_"
282+
if name[0].isdigit():
283+
name = f"x{name}"
284+
return name

0 commit comments

Comments
 (0)