TypeTree - use UnknownObject, move clean_name to TypeTreeNode, cache… · K0lb3/UnityPy@42e1320 · GitHub
Skip to content

Commit 42e1320

Browse files
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

‎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

0 commit comments

Comments
 (0)