Better type hints

Author: mos9527

README.md

@@ -13,15 +13,15 @@ Used in
 Example usage in `rla_transport_defines`
 ```python
 env = UnityPy.load(...)
-from generated.Sekai.Streaming import TransportDefine
-from generated import UTTCGen_AsInstance
+from generated import UTTCGen_AsInstance, UTTCGen_GetClass
 
 for reader in filter(lambda x: x.type == ClassIDType.MonoBehaviour, env.objects):
     name = reader.peek_name()
-    if name.startswith("TransportDefine"):
-        instance = UTTCGen_AsInstance(reader, "Sekai.Streaming.TransportDefine")
-        # ...or `instance = UTTCGen_AsInstance(reader)` is `m_Script` field is accessible
-        instance : TransportDefine
+    if name.startswith("TransportDefine"):        
+        from generated.Sekai.Streaming import TransportDefine
+        # ...or TransportDefine = UTTCGen_GetClass("Sekai.Streaming.TransportDefine")
+        # ...or TransportDefine = UTTCGen_GetClass(reader.read(check_read=False))
+        instance = UTTCGen_AsInstance(TransportDefine, reader)                
         print(f'"{name}":({instance.validBones}, {instance.validBlendShapes}),')
         # Possibly modify on the instance itself and saving it is also possible
         instance.validBones[0] = "BadBone"

UnityPyTypetreeCodegen/__main__.py

@@ -62,11 +62,11 @@
     [
         "# fmt: off",
         "# Auto-generated by https://github.com/mos9527/UnityPyTypetreeCodegen",
-        "" "from typing import List, Union, Optional, TypeVar",
+        "" "from typing import List, Union, Optional, TypeVar, Type",
         "from UnityPy.files.ObjectReader import ObjectReader",
         "from UnityPy.classes import *",
         "from UnityPy.classes.math import (ColorRGBA, Matrix3x4f, Matrix4x4f, Quaternionf, Vector2f, Vector3f, Vector4f, float3, float4,)",
-        '''T = TypeVar("T")
+        '''
 UTTCG_Classes = dict()
 def UTTCGen(fullname: str, typetree: dict):
     """dataclass-like decorator for typetree classess with nested type support
@@ -121,37 +121,40 @@ def __save(self):
         clazz.__init__ = __init__
         clazz.__repr__ = __repr__
         clazz.__typetree__ = typetree
+        clazz.__fullname__ = fullname
         clazz.save = __save
         UTTCG_Classes[fullname] = clazz
         return clazz
     return __inner
 
 # Helper functions
+def UTTCGen_GetClasss(src: MonoBehaviour | str) -> Type:
+    """Get the class definition from MonoBehaviour or a full type name."""
+    if isinstance(src, MonoBehaviour):
+        script = src.m_Script.read()
+        src = script.m_ClassName
+        if script.m_Namespace:
+            src = f"{script.m_Namespace}.{src}"    
+    return UTTCG_Classes.get(src, None)
 
-def UTTCGen_AsInstance(src: MonoBehaviour | ObjectReader, fullname: str = None) -> T:
+T = TypeVar("T")
+def UTTCGen_AsInstance(cls : Type[T], src: MonoBehaviour | ObjectReader) -> T:
     """Instantiate a class from the typetree definition and the raw data.
 
     In most cases, this is the function you want to use.
     It will read the typetree data from the MonoBehaviour instance and instantiate the class with the data.
 
     Args:
-        src (MonoBehaviour | ObjectReader): The MonoBehaviour instance or ObjectReader to read from.
-        fullname (str): The full name of the class to read. If None, it will be read from the MonoBehaviour instance's `m_Script` property.
+        cls: The class to instantiate. This should be a class that has been decorated with the UTTCGen decorator.
+        src (MonoBehaviour | ObjectReader): The MonoBehaviour instance or ObjectReader to read from.        
 
     Returns:
-        T: An instance of the class defined by the typetree.
+        An instance of the class defined by the typetree.
     """
-    if not fullname and isinstance(src, MonoBehaviour):
-        script = src.m_Script.read()
-        fullname = script.m_ClassName
-        if script.m_Namespace:
-            fullname = f"{script.m_Namespace}.{fullname}"    
-    clazz = UTTCG_Classes.get(fullname, None)
-    assert clazz is not None, f"Class definition for {fullname} not found"    
     if isinstance(src, MonoBehaviour):
         src = src.object_reader
-    raw_def = src.read_typetree(clazz.__typetree__, check_read=False)
-    instance = clazz(object_reader=src, **raw_def)
+    raw_def = src.read_typetree(cls.__typetree__, check_read=False)
+    instance = cls(object_reader=src, **raw_def)
     return instance
 ''',
     ]