refactor anm

lib/skn_anm.py

@@ -5,10 +5,27 @@ import math
 import prettyprinter
 from prettyprinter import pprint
 
-
 prettyprinter.install_extras()
 
 
+def with_offset(func):
+    def wrapped_do_offset(*args, **kwargs):
+        offset = kwargs.pop("offset", None)
+
+        b = args[1]
+        old = b.tell()
+        if offset:
+            b.seek(offset)
+
+        r = func(*args, **kwargs)
+
+        if offset:
+            b.seek(old)
+        return r
+
+    return wrapped_do_offset
+
+
 NAM_FlagHasHandles = (1 << 0)
 NAM_FlagCaseInsensitive = (1 << 1)
 ANM_NODE_FLAG_POS_ANIMATED = (1 << 0)
@@ -30,15 +47,11 @@ class Vec3:
         return math.sqrt(math.pow(self.x - other.x, 2) + math.pow(self.y - other.y, 2) + math.pow(self.z - other.z, 2))
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
         x = read_float_and_print("Vec3->x", b)
         y = read_float_and_print("Vec3->y", b)
         z = read_float_and_print("Vec3->z", b)
-        if offset:
-            b.seek(old)
         return cls(x=x, y=y, z=z)
 
 
@@ -47,16 +60,12 @@ class Vec3Indexed(Vec3):
     index: int
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
         x = read_float_and_print("Vec3->x", b)
         y = read_float_and_print("Vec3->y", b)
         z = read_float_and_print("Vec3->z", b)
         index = read_int_and_print("Vec3->index", b)
-        if offset:
-            b.seek(old)
         return cls(x=x, y=y, z=z, index=index)
 
 
@@ -65,19 +74,56 @@ class Vec4(Vec3):
     w: float
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
         x = read_float_and_print("Vec4->x", b)
         y = read_float_and_print("Vec4->y", b)
         z = read_float_and_print("Vec4->z", b)
         w = read_float_and_print("Vec4->w", b)
-        if offset:
-            b.seek(old)
         return cls(x=x, y=y, z=z, w=w)
 
 
+@dataclass
+class AnmKey:
+    time: int
+
+    @classmethod
+    @with_offset
+    def create(cls, b):
+        time = read_byte_and_print("ANM_Key->Time", b)
+        return cls(time=time)
+
+
+@dataclass
+class AnmKeys:
+    keys_type: int
+    ort_before: int
+    ort_after: int
+    flags: int
+    p_interpolate: int
+    num_keys: int
+    keys: list
+
+    @classmethod
+    @with_offset
+    def create(cls, b):
+        keys_type = read_byte_and_print("ANM_Keys->Type", b)
+        ort_before = read_byte_and_print("ANM_Keys->ORTBefore", b)
+        ort_after = read_byte_and_print("ANM_Keys->ORTAfter", b)
+        flags = read_byte_and_print("ANM_Keys->Flags", b)
+        p_interpolate = read_int_and_print("ANM_Keys->Interpolate", b)
+        num_keys = read_int_and_print("ANM_Keys->NumKeys", b)
+        p_keys = read_int_and_print("ANM_Keys->Keys", b)
+        keys = []
+        _old = b.tell()
+        b.seek(p_keys)
+        for _ in range(num_keys):
+            keys.append(AnmKey.create(b))
+        b.seek(_old)
+        return cls(keys_type=keys_type, ort_before=ort_before, ort_after=ort_after, flags=flags,
+                   p_interpolate=p_interpolate, num_keys=num_keys, keys=keys)
+
+
 @dataclass
 class AnmObj:
     custom: int = field(repr=False)
@@ -88,10 +134,8 @@ class AnmObj:
     clusters: List['AnmCluster']
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
         custom = read_int_and_print("ANM_Obj->Custom", b)
         skin = read_int_and_print("ANM_Obj->Skin", b)
         zero = read_int_and_print("ANM_Obj->ZeroCluster", b)
@@ -106,9 +150,8 @@ class AnmObj:
             clusters.append(cluster)
 
         b.seek(_old1)
-        if offset:
+        r = cls(custom=custom, skin=skin, zero_cluster=zero, total_number_vertices=total_vertices,
-            b.seek(old)
+                num_clusters=num_clusters, clusters=clusters)
-        r = cls(custom=custom, skin=skin, zero_cluster=zero, total_number_vertices=total_vertices, num_clusters=num_clusters, clusters=clusters)
         for cluster in clusters:
             cluster.obj = r
         return r
@@ -125,15 +168,13 @@ class AnmCluster:
     bone_name: str = ""
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
         custom = read_int_and_print("ANM_Cluster->Custom", b)
         obj_offset = read_int_and_print("ANM_Cluster->Obj", b)
         if obj_offset == 0:  # ugly hack, dunno why we need that
-            _ = read_int_and_print("ANM_Cluster->Custom", b)
+            custom = read_int_and_print("ANM_Cluster->Custom", b)
-            _ = read_int_and_print("ANM_Cluster->Obj", b)
+            obj_offset = read_int_and_print("ANM_Cluster->Obj", b)
         obj = None  # is set by parent Obj
         handle = read_int_and_print("ANM_Cluster->Handle", b)
         bbox1 = Vec3.create(b)
@@ -145,24 +186,21 @@ class AnmCluster:
             v = Vec3Indexed.create(b)
             vertices.append(v)
 
-        if offset:
+        return cls(custom=custom, obj=obj, handle=handle, bounding_box=bounding_box, num_vertices=num_vertices,
-            b.seek(old)
+                   vertices=vertices)
-        return cls(custom=custom, obj=obj, handle=handle, bounding_box=bounding_box, num_vertices=num_vertices, vertices=vertices)
 
 
 class NamDictionnary:
     @staticmethod
-    def create(b, cls, depth=1, offset=None):
+    @with_offset
+    def create(cls, b, depth=1):
         r = {}
 
-        old = b.tell()
-        if offset:
-            b.seek(offset)
-
         read_int_and_print("NAM_Dictionnary->pMemGroup", b)
         namemapper = read_int_and_print("NAM_Dictionnary->pNameMapper", b)
         flags = read_int_and_print("NAM_Dictionnary->Flags", b)
-        print("Flags: NAM_FlagHasHandles: %s, NAM_FlagCaseInsensitive: %s" % (flags & NAM_FlagHasHandles, flags & NAM_FlagCaseInsensitive))
+        print("Flags: NAM_FlagHasHandles: %s, NAM_FlagCaseInsensitive: %s" % (
+            flags & NAM_FlagHasHandles, flags & NAM_FlagCaseInsensitive))
         read_short_and_print("NAM_Dictionnary->NameHandle", b)
 
         b.seek(namemapper)
@@ -177,9 +215,7 @@ class NamDictionnary:
         read_int_and_print("TotalKeys", b)
         p_key_contexts = read_int_and_print("pKeyContexts", b)
 
-        print("MAP_Image->ANM_Skin->NAM_Dictionnary->MAP_Mapper->MAP_KeyContext")
         b.seek(p_key_contexts)
-        print("\tMAP_Image->ANM_Skin->NAM_Dictionnary->MAP_Mapper->MAP_KeyContext->MAP_BinSearchContext")
         _ = read_int_and_print("\tSize", b)
         read_int_and_print("\tpCallback", b)
         read_int_and_print("\tIndex", b)
@@ -193,7 +229,6 @@ class NamDictionnary:
             pindx = read_int(b)
             for _ in range(i):
                 pindx = read_int(b)
-            print("LUT INDEX: %s" % pindx)
             b.seek(pindx)
             if not flags & NAM_FlagHasHandles:
                 read_short_and_print("pPtr", b)
@@ -201,9 +236,9 @@ class NamDictionnary:
                 val = read_int_and_print("Value", b)
                 name = read_string_until_none_and_print("name", b)
                 if depth > 1:
-                    item = NamDictionnary.create(b, cls, depth-1, val)
+                    item = NamDictionnary.create(cls, b, depth - 1, offset=val)
                 else:
-                    item = cls.create(b, val)
+                    item = cls.create(b, offset=val)
                 r[name] = item
             else:
                 refcnt = read_short_and_print("RefCnt", b)
@@ -213,9 +248,6 @@ class NamDictionnary:
                 bone_name = read_string_until_none_and_print("BoneName", b)
                 r[str(i)] = {"BoneName": bone_name}
 
-        if offset:
-            b.seek(old)
-
         return r
 
 
@@ -224,22 +256,13 @@ class GeoMat4x4:
     m: List[float]
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
-
         m = []
         for _ in range(4):
             for _j in range(4):
                 m.append(read_float(b))
-
+        return cls(m=m)
-        r = cls(m=m)
-
-        if offset:
-            b.seek(old)
-
-        return r
 
 
 @dataclass
@@ -248,20 +271,11 @@ class GeoMat:
     s: GeoMat4x4
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
-
         mat_class = read_int(b)
         s = GeoMat4x4.create(b)
-
+        return cls(mat_class=mat_class, s=s)
-        r = cls(mat_class=mat_class, s=s)
-
-        if offset:
-            b.seek(old)
-
-        return r
 
 
 @dataclass
@@ -272,22 +286,13 @@ class GeoAffine:
     affine_class: int
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
-
         trans = Vec3.create(b)
         quat = Vec4.create(b)
         scale = Vec3.create(b)
         affine_class = read_int(b)
-
+        return cls(transform=trans, quat=quat, scale=scale, affine_class=affine_class)
-        r = cls(transform=trans, quat=quat, scale=scale, affine_class=affine_class)
-
-        if offset:
-            b.seek(old)
-
-        return r
 
 
 @dataclass
@@ -296,20 +301,11 @@ class AnmTransform:
     affine: GeoAffine
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
-
         time = read_int(b)
         affine = GeoAffine.create(b)
-
+        return cls(time=time, affine=affine)
-        r = cls(time=time, affine=affine)
-
-        if offset:
-            b.seek(old)
-
-        return r
 
 
 @dataclass
@@ -326,16 +322,13 @@ class AnmNode:
     stamp: int
     geo_mat: GeoMat
     transform: AnmTransform
-    pos_keys: int
+    pos_keys: Union[AnmKeys, int]
-    rot_keys: int
+    rot_keys: Union[AnmKeys, int]
-    scl_keys: int
+    scl_keys: Union[AnmKeys, int]
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
-
         custom = read_int_and_print("ANM_Node->Custom", b)
         node_type = read_short_and_print("ANM_Node->Type", b)
         x_form_type = read_short_and_print("ANM_Node->XformType", b)
@@ -352,28 +345,20 @@ class AnmNode:
         rot_keys = read_int_and_print("ANM_Node->pRotKeys", b)
         scl_keys = read_int_and_print("ANM_Node->pSclKeys", b)
 
-        """
+        if pos_keys > 0:
-        if nodeflags & ANM_NODE_FLAG_POS_ANIMATED:
+            pos_keys = AnmKeys.create(b, offset=pos_keys)
-            b.seek(posKeys)
-            read_byte_and_print("\t\tType", b)
-            read_byte_and_print("\t\tORTBefore", b)
-            read_byte_and_print("\t\tORTAfter", b)
-            read_byte_and_print("\t\tFlags", b)
-            read_int_and_print("\t\tpInterpolate", b)
-            read_int_and_print("\t\tNumberOfKeys", b)
-            read_int_and_print("\t\tpKeys", b)
-        """
 
-        r = cls(custom=custom, node_type=node_type, x_form_type=x_form_type, flags=flags, p_anim=p_anim,
+        if rot_keys > 0:
+            rot_keys = AnmKeys.create(b, offset=rot_keys)
+
+        if scl_keys > 0:
+            scl_keys = AnmKeys.create(b, offset=scl_keys)
+
+        return cls(custom=custom, node_type=node_type, x_form_type=x_form_type, flags=flags, p_anim=p_anim,
                    p_parent=p_parent, p_target=p_target, p_usercall=p_usercall, p_usercontext=p_usercontext,
                    stamp=stamp, geo_mat=geo_mat, transform=transform, pos_keys=pos_keys, rot_keys=rot_keys,
                    scl_keys=scl_keys)
 
-        if offset:
-            b.seek(old)
-
-        return r
-
 
 @dataclass
 class AnmAnim:
@@ -403,11 +388,8 @@ class AnmAnim:
     name: str
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
-
         custom = read_int_and_print("ANM_Anim->Custom", b)
         version = read_int_and_print("ANM_Anim->FileVersion", b)
         image = read_int_and_print("ANM_Anim->pImage", b)
@@ -430,20 +412,15 @@ class AnmAnim:
         otheranim = read_int_and_print("ANM_AnimExtra->OtherAnim", b)
         name = read_string_until_none(b)
 
-        treedict = NamDictionnary.create(b, AnmNode, depth=2, offset=treedict)
+        treedict = NamDictionnary.create(AnmNode, b, depth=2, offset=treedict)
-        defaultobjdict = NamDictionnary.create(b, AnmNode, offset=defaultobjdict)
+        defaultobjdict = NamDictionnary.create(AnmNode, b, offset=defaultobjdict)
 
-        r = cls(custom=custom, fileversion=version, image=image, memgroup=memgroup, start_time=start_time,
+        return cls(custom=custom, fileversion=version, image=image, memgroup=memgroup, start_time=start_time,
                    end_time=end_time, ticks_per_frame=ticks_per_frame, framerate=framerate, tree_dictionnary=treedict,
                    default_obj_dictionnary=defaultobjdict, num_nodes=num_nodes, extra_stamp=stamp, extra_flags=flags,
                    extra_time1=time1, extra_stamp1=stamp1, extra_time2=time2, extra_stamp2=stamp2, extra_factor=factor,
                    extra_other_anim=otheranim, name=name)
 
-        if offset:
-            b.seek(old)
-
-        return r
-
 
 @dataclass
 class AnmSkin:
@@ -468,11 +445,8 @@ class AnmSkin:
                 cluster.bone_name = nodedict[k]["BoneName"]
 
     @classmethod
-    def create(cls, b, offset=None):
+    @with_offset
-        old = b.tell()
+    def create(cls, b):
-        if offset:
-            b.seek(offset)
-
         custom = read_int_and_print("pCustom", b)
         version = read_int_and_print("FileVersion", b)
         image = read_int_and_print("pImage", b)
@@ -485,17 +459,14 @@ class AnmSkin:
         num_clusters = read_int_and_print("TotalNumberOfClusters", b)
         name = read_string_until_none_and_print("Name", b)
 
-        nodedict = NamDictionnary.create(b, AnmObj, offset=node)
+        nodedict = NamDictionnary.create(AnmObj, b, offset=node)
-        treedict = NamDictionnary.create(b, AnmObj, depth=2, offset=tree)
+        treedict = NamDictionnary.create(AnmObj, b, depth=2, offset=tree)
-        defaultobjdict = NamDictionnary.create(b, AnmObj, offset=defaultobj)
+        defaultobjdict = NamDictionnary.create(AnmObj, b, offset=defaultobj)
 
         # AnmSkin.link(nodedict, defaultobjdict)
-        r = cls(custom=custom, file_version=version, image=image, memgroup=memgroup, node_dictionnary=nodedict, tree_dictionnary=treedict, default_obj_dictionnary=defaultobjdict, total_number_trees=num_trees, total_number_objs=num_objs, total_number_clusters=num_clusters, name=name)
+        return cls(custom=custom, file_version=version, image=image, memgroup=memgroup, node_dictionnary=nodedict,
-
+                   tree_dictionnary=treedict, default_obj_dictionnary=defaultobjdict, total_number_trees=num_trees,
-        if offset:
+                   total_number_objs=num_objs, total_number_clusters=num_clusters, name=name)
-            b.seek(old)
-
-        return r
 
 
 @dataclass
@@ -514,4 +485,4 @@ class AnmFile:
         read_int_and_print("TotalSize", b)
         read_int_and_print("TotalBlocks", b)
         base = read_int_and_print("pBase", b)
-        return cls.create(b, base)
+        return cls.create(b, offset=base)