Add helper functions UTTCGen_AsInstance and LUT for generated classes

mos9527 · mos9527 · commit aa0aef92d849 · 2025-04-13T15:50:32.000+08:00
diff --git a/UnityPyTypetreeCodegen/__main__.py b/UnityPyTypetreeCodegen/__main__.py
@@ -51,6 +51,7 @@
     "Vector2": "Vector2f",
     "Vector3": "Vector3f",
     "Vector4": "Vector4f",
+    "Quaternion": "Quaternionf",
 }
 # XXX: Can't use attrs here since subclassing MonoBehavior and such - though defined by the typetree dump
 # seem to be only valid if the class isn't a property of another class
@@ -62,70 +63,116 @@
         "# fmt: off",
         "# Auto-generated by https://github.com/mos9527/UnityPyTypetreeCodegen",
         "" "from typing import List, Union, Optional, TypeVar",
+        "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")
-def UTTCGen(clazz : T) -> T:
-	"""dataclass-like decorator for typetree classess with nested type support
-	
-	limitations:
-	- the behavior is similar to slotted dataclasses where shared attributes are inherited
-	  but allows ommiting init of the parent if kwargs are not sufficient
-	- generally supports nested types, however untested and could be slow	
-    - and ofc, zero type checking and safeguards :/	
-	"""
-	RESERVED_KWS = {'object_reader'}
-    # Allow these to be propogated to the props
-	def __init__(self, **d):
-		def enusre_reserved(obj):
-			for k in RESERVED_KWS:
-				if hasattr(obj, k) and k in d:
-					setattr(obj, k, d[k])
-			return obj		
-		def reduce_init(clazz, **d):
-			types : dict = clazz.__annotations__
-			for k, sub in types.items():
-				if type(sub) == str:
-					sub = eval(sub) # attrs turns these into strings...why?
-				reduce_arg = getattr(sub, "__args__", [None])[0]
-				if isinstance(d[k], list):
-					if hasattr(reduce_arg, "__annotations__"):
-						setattr(self, k, [enusre_reserved(reduce_arg(**x)) for x in d[k]])
-					else:
-						setattr(self, k, [enusre_reserved(reduce_arg(x)) for x in d[k]])
-				elif isinstance(d[k], dict) and hasattr(sub, "__annotations__"):
-					setattr(self, k, enusre_reserved(sub(**d[k])))
-				else:
-					if isinstance(d[k], dict):
-						setattr(self, k, enusre_reserved(sub(**d[k])))
-					else:
-						setattr(self, k, enusre_reserved(sub(d[k])))			
-		for __base__ in clazz.__bases__:
-			types : dict = __base__.__annotations__
-			args = {k:d[k] for k in types if k in d}
-			if len(args) == len(types):
-				super(clazz, self).__init__(**args)
-				reduce_init(__base__, **d)
-				for k in args:
-					if not k in RESERVED_KWS: del d[k]
-		reduce_init(clazz, **d)
-		enusre_reserved(self)
-	def __repr__(self) -> str:
-		return f"{clazz.__name__}({', '.join([f'{k}={getattr(self, k)!r}' for k in self.__annotations__])})"
-	clazz.__init__ = __init__
-	clazz.__repr__ = __repr__
-	return clazz
+UTTCG_Classes = dict()
+def UTTCGen(fullname: str):
+    """dataclass-like decorator for typetree classess with nested type support
     
+    limitations:
+    - the behavior is similar to slotted dataclasses where shared attributes are inherited
+      but allows ommiting init of the parent if kwargs are not sufficient
+    - generally supports nested types, however untested and could be slow	
+    - and ofc, zero type checking and safeguards :/	
+    """    
+    def __inner(clazz: T) -> T:
+        RESERVED_KWS = {'object_reader'}
+        # Allow these to be propogated to the props
+        def __init__(self, **d):
+            def enusre_reserved(obj):
+                for k in RESERVED_KWS:
+                    if hasattr(obj, k) and k in d:
+                        setattr(obj, k, d[k])
+                return obj		
+            def reduce_init(clazz, **d):
+                types : dict = clazz.__annotations__
+                for k, sub in types.items():
+                    if type(sub) == str:
+                        sub = eval(sub) # attrs turns these into strings...why?
+                    reduce_arg = getattr(sub, "__args__", [None])[0]
+                    if isinstance(d[k], list):
+                        if hasattr(reduce_arg, "__annotations__"):
+                            setattr(self, k, [enusre_reserved(reduce_arg(**x)) for x in d[k]])
+                        else:
+                            setattr(self, k, [enusre_reserved(reduce_arg(x)) for x in d[k]])
+                    elif isinstance(d[k], dict) and hasattr(sub, "__annotations__"):
+                        setattr(self, k, enusre_reserved(sub(**d[k])))
+                    else:
+                        if isinstance(d[k], dict):
+                            setattr(self, k, enusre_reserved(sub(**d[k])))
+                        else:
+                            setattr(self, k, enusre_reserved(sub(d[k])))			
+            for __base__ in clazz.__bases__:
+                types : dict = __base__.__annotations__
+                args = {k:d[k] for k in types if k in d}
+                if len(args) == len(types):
+                    super(clazz, self).__init__(**args)
+                    reduce_init(__base__, **d)
+                    for k in args:
+                        if not k in RESERVED_KWS: del d[k]
+            reduce_init(clazz, **d)
+            enusre_reserved(self)
+        def __repr__(self) -> str:
+            return f"{clazz.__name__}({', '.join([f'{k}={getattr(self, k)!r}' for k in self.__annotations__])})"
+        clazz.__init__ = __init__
+        clazz.__repr__ = __repr__
+        UTTCG_Classes[fullname] = clazz
+        return clazz
+    return __inner
+
 # Helper functions
-def UTTCGen_Reread(Mono: MonoBehaviour):
-    script = Mono.m_Script.read()
-    fullname = script.m_ClassName
-    if script.m_Namespace:
-        fullname = f"{script.m_Namespace}.{fullname}"
-    definition = TYPETREE_DEFS.get(fullname, None)
+def UTTCGen_Reread(src: MonoBehaviour | ObjectReader, fullname: str = None) -> Tuple[dict, str]:
+    """Rereads the typetree data from the MonoBehaviour instance with our typetree definition.
+    
+    Args:
+        src: The MonoBehaviour instance or ObjectReader to read from.
+        fullname: The full name of the class to read. If None, it will be read from the MonoBehaviour instance's `m_Script` property.
+
+    Returns:
+        Tuple[dict, str]: [raw_def, fullname] where `raw_def` is the raw data read from the MonoBehaviour instance and `fullname` is the full name of the class.
+    """
+    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}"    
+    definition = UTTCG_Typetrees.get(fullname, None)
     assert definition is not None, f"Typetree definition for {fullname} not found"
-    raw_def = Mono.object_reader.read_typetree(definition)
+    if isinstance(src, MonoBehaviour):
+        src = src.object_reader
+    raw_def = src.read_typetree(definition, check_read=False)
     return raw_def, fullname
+
+def UTTCGen_Instantiate(raw_def : dict, fullname: str):
+    """Instantiate a class from the typetree definition and the raw data.
+
+    Args:
+        raw_def (dict): The raw data read from the MonoBehaviour instance.
+        fullname (str): The full name of the class to instantiate.
+    """
+    clazz = UTTCG_Classes.get(fullname, None)
+    assert clazz is not None, f"Class definition for {fullname} not found"    
+    instance = clazz(**raw_def)
+    return instance
+
+def UTTCGen_AsInstance(src: MonoBehaviour | ObjectReader, fullname: str = None) -> 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.
+
+    Returns:
+        T: An instance of the class defined by the typetree.
+    """
+    raw_def, fullname = UTTCGen_Reread(src, fullname)
+    instance = UTTCGen_Instantiate(raw_def, fullname)
+    return instance
 ''',
     ]
 )
@@ -286,6 +333,7 @@ def emit_line(*lines: str):
     logger.info(f"Subpass 2: Generating code for {namespace}")
     dp = defaultdict(lambda: -1)
     for clazz in topo:
+        fullname = f"{namespace}.{clazz}" if namespace else clazz
         fields = classname_nodes.get(clazz, None)
         if not fields:
             logger.debug(
@@ -297,7 +345,7 @@ def emit_line(*lines: str):
         clazz = translate_name(clazz)
         clazzes.append(clazz)
         clazz_fields = list()
-        emit_line(f"@UTTCGen")
+        emit_line(f"@UTTCGen('{fullname}')")
         if lvl1:
             parent = translate_type(fields[0].m_Type, strip=True, fallback=False)
             emit_line(f"class {clazz}({parent}):")
@@ -409,7 +457,7 @@ def __encoder(obj):
         return obj
 
     __open("__init__.py").write(
-        "\nTYPETREE_DEFS = " + json.dumps(fullname_nodes, indent=4, default=__encoder)
+        "\nUTTCG_Typetrees = " + json.dumps(fullname_nodes, indent=4, default=__encoder)
     )
 
 
