Properlly handle MonoBehaviour inheritance in generated code

Author: mos9527

UnityPyTypetreeCodegen/__main__.py

@@ -41,8 +41,7 @@
     "bool": "bool",
     "string": "str",
     "TypelessData": "bytes",
-    # -- Extra
-    "Array": "List[object]",
+    # -- Extra    
     "Byte[]": "bytes",
     "Byte": "int",
     "String": "str",
@@ -53,7 +52,13 @@
     "Vector3": "Vector3f",
     "Vector4": "Vector4f",
     "Quaternion": "Quaternionf",
+    "Object": "object",
+    "Array": "list",
+    "Type": "type",
+    "MethodInfo": "object",
+    "PropertyInfo": "object",
 }
+
 # 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
 # In which case the MonoBehavior attributes are inherited by the parent class and does not
@@ -304,11 +309,11 @@ def emit_line(*lines: str):
             f.write("\n")
 
     logger.info(
-        f"Subpass 1: Generating class dependency graph for {namespace or '<default namespace>'}"
+        f"Subpass 1: Generating class dependency graph for {namespace or "<default namespace>"}"
     )
     emit_line("# fmt: off")
     emit_line("# Auto-generated by https://github.com/mos9527/UnityPyTypetreeCodegen")
-    emit_line(f"# Python definition for {namespace or '<default namespace>'}", "")
+    emit_line(f"# Python definition for {namespace or "<default namespace>"}", "")
     if import_root:
         emit_line(f"from {import_root} import *")
     for clazz, parent in import_defs.items():
@@ -326,6 +331,7 @@ def emit_line(*lines: str):
     clazzes = list()
 
     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)
@@ -334,6 +340,8 @@ def emit_line(*lines: str):
                 f"Class {clazz} has no fields defined in TypeTree dump, skipped"
             )
             continue
+        lvl0 = list(filter(lambda field: field.m_Level == 0, fields))
+        lvl1 = list(filter(lambda field: field.m_Level == 1, fields))
         clazz = translate_name(clazz)
         clazzes.append(clazz)
         clazz_fields = list()
@@ -345,20 +353,47 @@ def __encoder(obj):
 
         clazz_typetree = json.dumps(fields, default=__encoder)
         emit_line(f"@UTTCGen('{fullname}', {clazz_typetree})")
-        emit_line(f"class {clazz}:")
-        for field in fields:
-            if field.m_Type == clazz:
-                continue
-            if field.m_Level > 1:
-                # Nested type. We should have already defined it
-                # with type hints
-                continue
-            name, type = field.m_Name, translate_type(
-                field.m_Type, typenames=classname_nodes | import_defs
-            )
-            emit_line(f"\t{declare_field(name, type, field.m_Type)}")
-            clazz_fields.append((name, type))
-
+        # Heuristic: If there is a lvl1 and a lvl0 field, it's a subclass
+        if lvl1 and lvl0:
+            parent = translate_type(fields[0].m_Type, strip=True, fallback=False)
+            emit_line(f"class {translate_name(clazz)}({translate_name(parent)}):")
+            if dp[parent] == -1:
+                # Reuse parent's fields with best possible effort
+                if pa_dep1 := getattr(UnityBuiltin, parent, None):
+                    dp[parent] = len(pa_dep1.__annotations__)
+                else:
+                    raise ValueError  # XXX: Should NEVER happen
+            pa_dep1 = dp[parent]
+            cur_dep1 = pa_dep1
+            for dep, (i, field) in enumerate(
+                filter(lambda field: field[1].m_Level == 1, enumerate(fields))
+            ):
+                if field.m_Level > 1:
+                    continue # Nested                
+                if dep < pa_dep1:
+                    # Skip parent fields at lvl1
+                    continue
+                if i + 1 < len(fields) and fields[i + 1].m_Type == "Array":
+                    field.m_Type = fields[i + 3].m_Type + "[]"
+                name, type = field.m_Name, translate_type(
+                    field.m_Type, typenames=classname_nodes | import_defs
+                )
+                emit_line(f"\t{declare_field(name, type, field.m_Type)}")
+                clazz_fields.append((name, type, field.m_Type))
+                cur_dep1 += 1
+            dp[clazz] = cur_dep1
+        else:
+            # No inheritance
+            emit_line(f"class {clazz}:")
+            for field in fields:
+                if field.m_Level > 1:
+                    continue # Nested
+                name, type = field.m_Name, translate_type(
+                    field.m_Type, typenames=classname_nodes | import_defs
+                )
+                emit_line(f"\t{declare_field(name, type, field.m_Type)}")
+                clazz_fields.append((name, type))
+            dp[clazz] = len(fields)
         if not clazz_fields:
             # Empty class. Consider MRO
             emit_line("\tpass")
@@ -437,11 +472,6 @@ def __main__():
         help="Unity version to use for typetree generation",
         default="2022.3.21f1",
     )
-    parser.add_argument(
-        "--backend",
-        help="Backend to use for code generation",
-        default="AssetStudio",
-    )
     parser.add_argument(
         "--filter",
         help="Filter classnames by regex",
@@ -481,13 +511,14 @@ def __main__():
     shutil.rmtree(args.outdir, ignore_errors=True)
     os.makedirs(args.outdir, exist_ok=True)
     typetree = dict()
-    gen = TypeTreeGenerator(args.unity_version, args.backend)
+    gen = TypeTreeGenerator(args.unity_version, "AssetStudio")
     def populate_gen():
         # https://github.com/UnityPy-Org/TypeTreeGeneratorAPI/pull/1
         for asm,clz in gen.get_class_definitions():
             try:
                 node = gen.get_nodes_as_json(asm, clz)
                 node = json.loads(node)
+                # https://github.com/UnityPy-Org/TypeTreeGeneratorAPI/pull/4                
                 typetree[clz] = node
             except Exception as e:
                 logger.warning(f"Skipping nodes for {asm}.{clz}: {e}")