UnityPyTypetreeCodegen/UnityPyTypetreeCodegen/__main__.py at main · mos9527/UnityPyTypetreeCodegen

This repository was archived by the owner on Feb 6, 2026. It is now read-only.
658 lines (610 loc) · 22.7 KB
"""Simple codegen for typetree classes
- Supports nested types
- Supports inheritance
- Automatically resolves import order and dependencies
- Generated classes support nested types, and will be initialized with the correct types even with nested dicts
- Cannot resolve namespace conflicts if the same class name is defined in multiple namespaces
- Missing type definitions are marked with # XXX: Fallback of {org_type} and typedefed as object
- Circular inheritance is checked and raises RecursionError
- The output order (imports, classes) are deterministic and lexicographically sorted
- The output is emitted in lieu of the Namespace structure of the TypeTree dump, presented as Python modules in directories
    python typetree_codegen.py <typetree_dump.json> <output_dir>
# From https://github.com/K0lb3/UnityPy/blob/master/generators/ClassesGenerator.py
BASE_TYPE_MAP = {
    "char": "str",
    "short": "int",
    "int": "int",
    "long long": "int",
    "unsigned short": "int",
    "unsigned int": "int",
    "unsigned long long": "int",
    "UInt8": "int",
    "UInt16": "int",
    "UInt32": "int",
    "UInt64": "int",
    "SInt8": "int",
    "SInt16": "int",
    "SInt32": "int",
    "SInt64": "int",
    "Type*": "int",
    "FileSize": "int",
    "float": "float",
    "double": "float",
    "bool": "bool",
    "string": "str",
    "TypelessData": "bytes",
    # -- Extra    
    "Byte[]": "bytes",
    "Byte": "int",
    "String": "str",
    "Int32": "int",
    "Single": "float",
    "Color": "ColorRGBA",
    "Vector2": "Vector2f",
    "Vector3": "Vector3f",
    "Vector4": "Vector4f",
    "Quaternion": "Quaternionf",
    "Object": "object",    
    "Type": "type",
    "MethodInfo": "object",
    "PropertyInfo": "object",    
Typetree_MonoBehaviour = [
      "m_Type": "MonoBehaviour",
      "m_Name": "Base",
      "m_MetaFlag": 0,
      "m_Level": 0
      "m_Type": "PPtr<GameObject>",
      "m_Name": "m_GameObject",
      "m_MetaFlag": 0,
      "m_Level": 1
      "m_Type": "int",
      "m_Name": "m_FileID",
      "m_MetaFlag": 0,
      "m_Level": 2
      "m_Type": "SInt64",
      "m_Name": "m_PathID",
      "m_MetaFlag": 0,
      "m_Level": 2
      "m_Type": "UInt8",
      "m_Name": "m_Enabled",
      "m_MetaFlag": 16384,
      "m_Level": 1
      "m_Type": "PPtr<MonoScript>",
      "m_Name": "m_Script",
      "m_MetaFlag": 0,
      "m_Level": 1
      "m_Type": "int",
      "m_Name": "m_FileID",
      "m_MetaFlag": 0,
      "m_Level": 2
      "m_Type": "SInt64",
      "m_Name": "m_PathID",
      "m_MetaFlag": 0,
      "m_Level": 2
      "m_Type": "string",
      "m_Name": "m_Name",
      "m_MetaFlag": 0,
      "m_Level": 1
      "m_Type": "Array",
      "m_Name": "Array",
      "m_MetaFlag": 16384,
      "m_Level": 2
      "m_Type": "int",
      "m_Name": "size",
      "m_MetaFlag": 0,
      "m_Level": 3
      "m_Type": "char",
      "m_Name": "data",
      "m_MetaFlag": 0,
      "m_Level": 3
# 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
# initialize the property class
# XXX: Need some boilerplate to handle this
HEADER = "\n".join(
        "# fmt: off",
        "# Auto-generated by https://github.com/mos9527/UnityPyTypetreeCodegen",
        "" "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,)",
UTTCG_Classes = dict()
def UTTCGen(fullname: str, typetree: List[dict]):
    """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 :/	
    """    
    REFERENCED_ARGS = {'object_reader'}
    def __inner(clazz: T) -> T:
        # Allow these to be propogated to the props
        def __init__(self, **d):        
            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?
                    while sub.__name__ == "Optional":
                        sub = sub.__args__[0]  # Reduce Optional[T] -> T								
                    reduce_arg = getattr(sub, "__args__", [None])[0]					
                    if k in REFERENCED_ARGS: # Directly refcounted
                        reduce_arg = sub = lambda x: x                         
                    if reduce_arg is not None and isinstance(d[k], list):
                        if hasattr(reduce_arg, "__annotations__") or hasattr(reduce_arg, "__args__"):
                            setattr(self, k, [reduce_arg(**x) for x in d[k]])
                            setattr(self, k, [reduce_arg(x) for x in d[k]])
                    elif reduce_arg is not None and isinstance(d[k], dict) and hasattr(sub, "__annotations__"):
                        setattr(self, k, sub(**d[k]))
                        # Reduce typings to respective types
                        if hasattr(sub, "__origin__") and sub.__origin__ is not None:
	                        sub = sub.__origin__		
                        if isinstance(d[k], dict):
                            # Special cases
                            # Color with `rgba` field
                            if sub == ColorRGBA and 'rgba' in d[k]:
                                rgba = d[k]['rgba']
                                r = ((rgba >> 24) & 0xFF) / 255
                                g = ((rgba >> 16) & 0xFF) / 255
                                b = ((rgba >> 8) & 0xFF) / 255
                                setattr(self, k, sub(r, g, b, a))
                                setattr(self, k, sub(**d[k]))
                            setattr(self, k, sub(d[k]))
            def reduce_base(clazz, **d):	
                for __base__ in clazz.__bases__:
                    if hasattr(__base__, "__annotations__"):
                        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)                       
                    reduce_base(__base__, **d)
            reduce_base(clazz, **d)               
            reduce_init(clazz, **d)            
        def __repr__(self) -> str:
            return f"{clazz.__name__}({', '.join([f'{k}={getattr(self, k)!r}' for k in self.__annotations__])})"
        def __save(self):
            self.object_reader.save_typetree(self, self.__typetree__)
        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_GetClass(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)
T = TypeVar("T")
def UTTCGen_AsInstance(cls : Type[T], src: MonoBehaviour | ObjectReader, check_read : bool = True) -> 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.
        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.        
        check_read (bool): Whether to check if all fields are read. Defaults to True.
    Returns:
        An instance of the class defined by the typetree.
    if isinstance(src, MonoBehaviour):
        src = src.object_reader
    raw_def = src.read_typetree(cls.__typetree__, check_read=check_read)
    instance = cls(object_reader=src, **raw_def)
    return instance
from collections import defaultdict
import argparse, json
def translate_name(m_Name: str, **kwargs):
    """A la https://github.com/K0lb3/UnityPy/blob/b811a2942297b5d8107e9a10249df80a87492282/UnityPy/helpers/TypeTreeNode.py#L361.
    With extra handling for templated/generic types and reserved keywords.
    NG = "<>|`="
    m_Name = m_Name.replace("<>", "__generic_")  # Generic templates
    m_Name = m_Name.replace("<", "_").replace(">", "_")  # Templated
    for c in NG:
        m_Name = m_Name.replace(c, "_")
    RESERVED_NAMES = {
        "class",
        "def",
        "return",
        "if",
        "else",
        "elif",
        "for",
        "while",
        "in",
        "is",
        "not",
        "and",
        "or",
        "from",
        "import",
        "as",
        "with",
        "try",
        "except",
        "finally",
        "raise",
        "assert",
        "break",
        "continue",
        "pass",
        "yield",
        "True",
        "False",
    if m_Name in RESERVED_NAMES:
        m_Name = "_" + m_Name
    return m_Name
from UnityPy import classes as UnityBuiltin
from TypeTreeGeneratorAPI import TypeTreeNode
from logging import getLogger
import os, shutil
from typing import Dict, List
logger = getLogger("codegen")
def translate_type(
    m_Type: str, strip=False, fallback=True, typenames: dict = dict(), feild_nexts : list = [], parent_name = None
    if m_Type == parent_name:
        return 'object' # XXX Recusrive. Python doesn't like those
    if m_Type in BASE_TYPE_MAP:
        return BASE_TYPE_MAP[m_Type]
    if getattr(UnityBuiltin, m_Type, None):
        return m_Type
    if m_Type in typenames:
        return m_Type
    if m_Type.endswith("[]") or m_Type == 'List`1':        
        if feild_nexts:
            m_Type = translate_type(feild_nexts[3].m_Type, strip, fallback, typenames, feild_nexts, parent_name)
        elif m_Type.endswith("[]") :
            m_Type = m_Type[:-2]
        else:
            m_Type = None
        if not strip:
            if m_Type:
                return f"List[{m_Type}]"
            else:
                logger.warning(f"Unknown list type of element {m_Type}, using fallback")
                return 'list'
        else:
            return m_Type
    if m_Type.startswith("PPtr<"):
        m_Type = translate_type(m_Type[5:-1], strip, fallback, typenames, feild_nexts, parent_name)
        if not strip:
            return f"PPtr[{m_Type}]"
        else:
            return m_Type
    if fallback:
        logger.warning(f"Unknown type {m_Type}, using fallback")
        return "object"
        return m_Type
def declare_field(name: str, type: str, org_type: str = None):
    name = translate_name(name)
    if type not in {"object", "List[object]", "PPtr[object]"}:
        return f"{name} : {type}"
        # We'd skip parsing these further if we don't know the type
        if type == "object":
            type = 'dict'
        if type == "List[object]":
            type = 'List[dict]'
        return f"{name} : {type} # XXX: Fallback of {org_type}"
from io import TextIOWrapper
def topsort(graph: dict):
    # Sort the keys in topological order
    # We don't assume the guarantee otherwise
    graph = {k: list(sorted((i for i in v if i != k))) for k, v in graph.items()}
    vis = defaultdict(lambda: 0)
    topo = list()
    def dfs(u):
        vis[u] = 1
        for v in graph.get(u, []):
            if vis[v] == 1:
                return False
            if vis[v] == 0 and not dfs(v):
                return False
        vis[u] = 2
        topo.append(u)
        return True
    flag = 1
    for clazz in graph:
        if not vis[clazz]:
            flag &= dfs(clazz)
    # XXX: Shouldn't happen. Need to figure out how this is possible
    # assert flag, "graph contains cycle"
    return topo
def process_namespace(
    f: TextIOWrapper,
    classname_nodes: Dict[str, List[TypeTreeNode]],
    namespace: str = "",
    import_root: str = "",
    import_defs: dict = dict(),
    def emit_line(*lines: str):
        for line in lines:
            f.write(line)
            f.write("\n")
        if not lines:
            f.write("\n")
    logger.info(
        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>"}", "")
    if import_root:
        emit_line(f"from {import_root} import *")
    for clazz, parent in import_defs.items():
        emit_line(f"from {import_root}{parent or ''} import {clazz}")
    emit_line()
    # Emit by topo order
    graph = {
        clazz: {
            translate_type(field.m_Type, strip=True, fallback=False, feild_nexts=fields[i:]) for (i,field) in enumerate(fields)
        for clazz, fields in classname_nodes.items()
    topo = topsort(graph)
    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)
        if not fields:
            logger.debug(
                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()
        def __encoder(obj):
            if isinstance(obj, TypeTreeNode):
                return obj.__dict__
            return obj
        clazz_typetree = json.dumps(fields, default=__encoder)
        emit_line(f"@UTTCGen('{fullname}', {clazz_typetree})")
        # 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, feild_nexts=fields[0:], parent_name=clazz)
            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
                name, type = field.m_Name, translate_type(
                    field.m_Type, typenames=classname_nodes | import_defs, feild_nexts=fields[i:], parent_name=clazz
                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 i, field in enumerate(fields):
                if field.m_Level > 1:
                    continue # Nested
                name, type = field.m_Name, translate_type(
                    field.m_Type, typenames=classname_nodes | import_defs, feild_nexts=fields[i:], parent_name=clazz
                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")
def process_typetree(fullname_nodes: Dict[str, List[TypeTreeNode]], outdir: str):
    handles: Dict[str, TextIOWrapper] = dict()
    def __open(fname: str):
        fname = os.path.join(outdir, fname)
        if fname not in handles:
            os.makedirs(os.path.dirname(fname), exist_ok=True)
            handles[fname] = open(fname, "w")
        return handles[fname]
    namespaces = defaultdict(dict)
    namespacesT = defaultdict(None)
    logger.info("Pass 1: Building namespace")
    for key in fullname_nodes:
        fullkey = key.split(".")
        if len(fullkey) == 1:
            namespace, clazz = None, fullkey[0]
        else:
            namespace, clazz = fullkey[:-1], fullkey[-1]
            namespace = ".".join(namespace)
        namespaces[namespace][clazz] = fullname_nodes[key]
        if clazz not in namespacesT:
            namespacesT[clazz] = namespace
        else:
            logger.error(
                f"Class {clazz} already defined in {namespacesT[clazz]} but found again in {namespace}"
            logger.error(
                f"Need manual intervention to resolve the conflict. Using first definition for now."
    logger.info("Pass 2: Generating import graph")
    # Build import graph
    namespaceDeps = defaultdict(set)
    for namespace, classname_nodes in namespaces.items():
        for clazz, fields in classname_nodes.items():
            for i, field in enumerate(fields):
                if type(field) != TypeTreeNode:
                    field = fields[i] = TypeTreeNode(**field)
                ftype = translate_type(field.m_Type, strip=True, fallback=False)
                if ftype in namespacesT and namespacesT[ftype] != namespace:
                    namespaceDeps[namespace].add(ftype)
    logger.info("Pass 3: Emitting namespace as Python modules")
    __open("__init__.py").write(HEADER)
    # XXX: This part can be trivally parallelized
    for namespace, classname_nodes in sorted(
        namespaces.items(), key=lambda x: x[0].count(".") if x[0] else 0
        # CubismTaskHandler -> generated/__init__.py
        # Live2D.Cubism.Core.CubismMoc -> generated/Live2D/Cubism/Core/__init__.py
        if namespace:
            ndots = namespace.count(".") + 2
            dotss = "." * ndots
            f = __open(os.path.join(*namespace.split("."), "__init__.py"))
            deps = {k: namespacesT[k] for k in namespaceDeps[namespace]}
            deps = dict(sorted(deps.items()))
            process_namespace(f, classname_nodes, namespace, dotss, deps)
        else:
            f = __open("__init__.py")
            process_namespace(f, classname_nodes, namespace)
import re, logging
from TypeTreeGeneratorAPI import TypeTreeGenerator
def __main__():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--unity-version",
        help="Unity version to use for typetree generation",
        default="2022.3.21f1",
    parser.add_argument(
        "--filter",
        help="Filter classnames by regex",
        default=".*",
    parser.add_argument(
        "--json",
        help="[JSON] Load tree dump in json format {str[fullname]: List[TypeTreeNode]},...",
    parser.add_argument(
        "--asm-dir",
        help="[Asm] Load typetree dump from game assembly DLL folder",
        type=str
    parser.add_argument(
        "--il2cpp",
        help="[IL2CPP] Load typetree dump from IL2CPP binaries",
        type=str
    parser.add_argument(
        "--metadata",
        help="[IL2CPP] Load typetree dump from metadata files",
        type=str
    parser.add_argument(
        "--log-level",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        default="WARNING",
    parser.add_argument(
        "--outdir",
        help="Output directory for generated code",
        default="generated",
    args = parser.parse_args()
    logging.basicConfig(level=args.log_level)
    shutil.rmtree(args.outdir, ignore_errors=True)
    os.makedirs(args.outdir, exist_ok=True)
    typetree = dict()
    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
                if node and node[0]["m_Level"] != 0:
                    node = Typetree_MonoBehaviour + node
                typetree[clz] = node
            except Exception as e:
                logger.warning(f"Skipping nodes for {asm}.{clz}: {e}")
    if args.asm_dir:
        print("Loading .NET Assemblies", args.asm_dir)
        for dll in os.listdir(args.asm_dir):
            if not dll.lower().endswith(".dll"):
                continue
            try:
                print(f"Loading {dll}")
                gen.load_dll(open(os.path.join(args.asm_dir, dll), "rb").read())
            except Exception as e:
                logger.warning(f"Skipping {dll}: {e}")
        populate_gen()
    elif args.il2cpp and args.metadata:
        print("Loading IL2CPP", args.il2cpp, args.metadata)
        gen.load_il2cpp(open(args.il2cpp, "rb").read(), open(args.metadata, "rb").read())
        populate_gen()
    elif args.json:
        print("Loading JSON", args.json)
        with open(args.json, "r") as f:
            typetree = json.load(f)
        raise ValueError("No valid input source specified.")
    if typetree:
        with open(".typetree.json", "w") as f:
            json.dump(typetree, f, indent=4)
        regex = re.compile(args.filter)
        typetree = {k: v for k, v in typetree.items() if regex.match(k)}
        process_typetree(typetree, args.outdir)
        return 0
    return -1
if __name__ == "__main__":
    sys.exit(__main__())
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

__main__.py

Latest commit

History

__main__.py

File metadata and controls

main.py

main.py
