diff options
Diffstat (limited to 'libunity.py')
| -rw-r--r-- | libunity.py | 601 |
1 files changed, 294 insertions, 307 deletions
diff --git a/libunity.py b/libunity.py index 5559c5f..097e8b8 100644 --- a/libunity.py +++ b/libunity.py @@ -1,190 +1,11 @@ #!/usr/bin/env python3 +import copy import enum +import sys # python3 -m pip install pyyaml import yaml -def getObjectIds(fx): - fx_ids = set() - with open(fx, "r") as f: - for line in f: - parts = line.split() - if line.startswith("---") and len(parts) == 3: - obj_id = parts[2] - # Remove & from beginning of object ID. - obj_id = obj_id[1:] - fx_ids.add(obj_id) - return fx_ids - -def getOldIds(old_fx_ids, unified_fx_ids): - id_mapping = {} - - for old_id in old_fx_ids: - id_mapping[old_id] = str(old_id) - unified_fx_ids.add(int(old_id)) - - return id_mapping - -# Map an object ID from old namespace to unified. -def getNewIds(old_fx_ids, unified_fx_ids): - id_mapping = {} - - for old_id in old_fx_ids: - new_id = int(old_id) - - # 9100000 is the ID of the animator. There's only one, and we take care - # to merge it. So no need to create a new identifier for it. - if new_id != 9100000: - while (new_id in unified_fx_ids or new_id in old_fx_ids): - new_id += 1 - - unified_fx_ids.add(new_id) - id_mapping[old_id] = str(new_id) - - return id_mapping - -def replaceIds(old_str, id_map): - result = old_str - lines = [] - for line in result.split("\n"): - for old_id, new_id in id_map.items(): - line = line.replace(old_id, new_id) - lines.append(line) - return '\n'.join(lines) - -def cat(path): - lines = [] - with open(path, "r") as f: - for line in f: - lines.append(line) - return ''.join(lines) - -# Extract the first block whose '---' line matches `prefix`. -def extractFirstBlock(fx, begin_prefix, end_prefix): - lines = [] - in_block = False - for line in fx.split("\n"): - if in_block and line.startswith(end_prefix): - break - if line.startswith(begin_prefix): - in_block = True - if in_block: - lines.append(line) - return '\n'.join(lines) + "\n" - -def mergeAnimatorControllers(fx0, fx1, fx0_ids, fx1_ids): - fx0_anim = extractFirstBlock(cat(fx0), "--- !u!91", "---") - fx1_anim = extractFirstBlock(cat(fx1), "--- !u!91", "---") - - fx0_anim_params = extractFirstBlock(fx0_anim, " m_AnimatorParameters:", " m_AnimatorLayers:") - fx0_anim_params = "\n".join(fx0_anim_params.split("\n")[1:]) - fx0_anim_params = replaceIds(fx0_anim_params, fx0_ids) - - fx0_anim_layers = extractFirstBlock(fx0_anim, " m_AnimatorLayers:", "---") - fx0_anim_layers = "\n".join(fx0_anim_layers.split("\n")[1:]) - fx0_anim_layers = replaceIds(fx0_anim_layers, fx0_ids) - - fx1_anim_params = extractFirstBlock(fx1_anim, " m_AnimatorParameters:", " m_AnimatorLayers:") - fx1_anim_params = "\n".join(fx1_anim_params.split("\n")[1:]) - fx1_anim_params = replaceIds(fx1_anim_params, fx1_ids) - - fx1_anim_layers = extractFirstBlock(fx1_anim, " m_AnimatorLayers:", "---") - fx1_anim_layers = "\n".join(fx1_anim_layers.split("\n")[1:]) - fx1_anim_layers = replaceIds(fx1_anim_layers, fx1_ids) - - result = """ ---- !u!91 &9100000 -AnimatorController: - m_ObjectHideFlags: 0 - m_CorrespondingSourceObject: {fileID: 0} - m_PrefabInstance: {fileID: 0} - m_PrefabAsset: {fileID: 0} - m_Name: FX - serializedVersion: 5 - m_AnimatorParameters: -"""[1:] - - result += fx0_anim_params - result += fx1_anim_params - - result += " m_AnimatorLayers:\n" - - result += fx0_anim_layers - result += fx1_anim_layers - - return result - -# Merge two FX layers. -# fx0, fx1 are both paths to animators. -# The object IDs for fx1 are reassigned to not collide with those used by fx0. -def mergeFX(fx0, fx1): - fx0_ids = getObjectIds(fx0) - fx1_ids = getObjectIds(fx1) - - # Get unique identifiers for all objects in both layers. - new_ids = set() - fx0_id_mapping = getOldIds(fx0_ids, new_ids) - fx1_id_mapping = getNewIds(fx1_ids, new_ids) - - # Merge animators - anim_ctrl = mergeAnimatorControllers(fx0, fx1, fx0_id_mapping, fx1_id_mapping) - - # Remove animators and prefix - fx0_str = cat(fx0) - fx1_str = cat(fx1) - - fx0_anim = extractFirstBlock(fx0_str, "--- !u!91", "---") - fx1_anim = extractFirstBlock(fx1_str, "--- !u!91", "---") - - fx0_str = fx0_str.replace(fx0_anim, "") - fx1_str = fx1_str.replace(fx1_anim, "") - - prefix = """ -%YAML 1.1 -%TAG !u! tag:unity3d.com,2011: -"""[1:] - - fx0_str = fx0_str.replace(prefix, "") - fx1_str = fx1_str.replace(prefix, "") - - # Replace old IDs - #fx0_str = replaceIds(fx0_str, fx0_id_mapping) - fx1_str = replaceIds(fx1_str, fx1_id_mapping) - - # Output - return deleteEmptyLines(prefix + anim_ctrl + fx0_str + fx1_str) - -def deleteEmptyLines(fx): - return fx.replace("\n\n", "\n") - -# Import a Unity YAML, modifying it to make it legal YAML 1.1. -def importUnityYaml(fx_old, fx_new): - lines = [] - with open(fx_old, "r") as f: - first_document = True - for line in f: - # Add end-of-document indicators. - if line.startswith("---"): - if not first_document: - lines.append("...\n") - first_document = False - - # Remove class ID tag from each block. - if line.startswith("---"): - parts = line.split() - lines.append(parts[0] + " " + parts[2] + "\n") - continue - lines.append(line) - - lines.append("...\n") - with open(fx_new, "w") as f: - for line in lines: - f.write(f"{line}") - -def AnimatorControllerConstructor(loader, tag_suffix, node): - print("Got animator node: {}".format(node)) - return None - class Node: def __init__(self): # Optional. In Unity, this is the fileID of an object. Not all YAML @@ -248,6 +69,10 @@ class Sequence(Node): return child + def forEach(self, cb): + for k in self.sequence: + cb(k) + class Mapping(Node): def __init__(self): super().__init__() @@ -307,33 +132,178 @@ class Mapping(Node): return child + def forEach(self, cb): + for k, v in self.mapping.items(): + cb(v) + +def classId(anchor): + # AnimatorController + if anchor.startswith("91"): + return "91" + # MonoBehaviour + if anchor.startswith("114"): + return "114" + # BlendTree + if anchor.startswith("206"): + return "206" + # AnimatorStateTransition + if anchor.startswith("1101"): + return "1101" + # AnimatorState + if anchor.startswith("1102"): + return "1102" + # AnimatorStateMachine + if anchor.startswith("1107"): + return "1107" + # AnimatorTransition + if anchor.startswith("1109"): + return "1109" + + # IDK what this is lmao + if anchor.startswith("74"): + return "74" + if anchor.startswith("115"): + return "115" + + raise Exception("Unrecognized object: {}".format(anchor)) + class UnityDocument(Mapping): def __str__(self): return super().__str__() def classId(self): - # AnimatorController - if self.anchor.startswith("91"): - return "91" - # MonoBehaviour - if self.anchor.startswith("114"): - return "114" - # BlendTree - if self.anchor.startswith("206"): - return "206" - # AnimatorStateTransition - if self.anchor.startswith("1101"): - return "1101" - # AnimatorState - if self.anchor.startswith("1102"): - return "1102" - # AnimatorStateMachine - if self.anchor.startswith("1107"): - return "1107" - # AnimatorTransition - if self.anchor.startswith("1109"): - return "1109" - raise Exception("Unrecognized object: {}".format(self.anchor)) + return classId(self.anchor) + +class UnityAnimator(): + def __init__(self): + self.nodes = [] + self.id_to_node = {} + + def addNodes(self, nodes): + for node in nodes: + self.nodes.append(node) + if node.anchor == None: + raise Exception("Node is missing anchor: {}".format(str(node))) + if node.anchor in self.id_to_node: + raise Exception("Duplicate anchor: {}, node 1: {}, node 2: {}".format(node.anchor, str(node), str(self.id_to_node[node.anchor]))) + self.id_to_node[node.anchor] = node + + def getUniqueId(self, anchor): + if anchor in self.id_mapping.keys(): + return self.id_mapping[anchor] + + if classId(anchor) in self.class_to_next_id: + new_id = self.class_to_next_id[classId(anchor)] + self.class_to_next_id[classId(anchor)] += 1 + self.id_mapping[anchor] = new_id + else: + new_id = int("%s%05d" % (classId(anchor), 0)) + next_id = new_id + 1 + self.class_to_next_id[classId(anchor)] = next_id + self.id_mapping[anchor] = new_id + + #print("Map {} to {}".format(anchor, self.id_mapping[anchor]), file=sys.stderr) + return self.id_mapping[anchor] + + def mergeIterator(self, v): + if hasattr(v, "mapping"): + # Don't relabel anything that's defined in an external file. + # TODO(yum) do this. + if 'fileID' in v.mapping and not 'guid' in v.mapping: + if v.mapping['fileID'] != '0': + old_id = v.mapping['fileID'] + new_id = self.getUniqueId(old_id) + v.mapping['fileID'] = str(new_id) + if hasattr(v, "forEach"): + v.forEach(self.mergeIterator) + + def popNodeOfClass(self, classId): + result = None + for node in self.nodes: + if node.classId() == classId: + result = node + self.nodes.remove(result) + break + del self.id_to_node[result.anchor] + return result + + def pushNode(self, node): + self.nodes.append(node) + self.id_to_node[node.anchor] = node + + def mergeAnimatorControllers(self, ctrl0, ctrl1): + ctrl0 = copy.deepcopy(ctrl0) + ctrl1 = copy.deepcopy(ctrl1) + + self.id_mapping0 = {} + self.id_mapping1 = {} + + p0 = ctrl0.mapping['AnimatorController'].mapping['m_AnimatorParameters'] + p1 = ctrl1.mapping['AnimatorController'].mapping['m_AnimatorParameters'] + + a0 = ctrl0.mapping['AnimatorController'].mapping['m_AnimatorLayers'] + a1 = ctrl1.mapping['AnimatorController'].mapping['m_AnimatorLayers'] + + self.id_mapping = self.id_mapping0 + p0.forEach(self.mergeIterator) + a0.forEach(self.mergeIterator) + + # Hack to prevent ctrl1 from getting a new ID for the animator. + del self.class_to_next_id['91'] + + self.id_mapping = self.id_mapping1 + p1.forEach(self.mergeIterator) + a1.forEach(self.mergeIterator) + + p0.sequence += p1.sequence + a0.sequence += a1.sequence + + return ctrl0 + + def merge(self, other): + # Mapping from class ID (string) to next ID (int) + self.class_to_next_id = {} + + ctrl0 = self.popNodeOfClass('91') + ctrl1 = other.popNodeOfClass('91') + merged_anim = self.mergeAnimatorControllers(ctrl0, ctrl1) + + # Mapping from class ID (string) to new class ID (int) + self.id_mapping = self.id_mapping0 + for node in self.nodes: + new_id = self.getUniqueId(node.anchor) + node.anchor = str(new_id) + node.forEach(self.mergeIterator) + + self.id_mapping = self.id_mapping1 + for node in other.nodes: + new_id = self.getUniqueId(node.anchor) + node.anchor = str(new_id) + node.forEach(self.mergeIterator) + + nodes = self.nodes + self.nodes = [] + self.id_to_node = {} + self.pushNode(merged_anim) + self.addNodes(nodes) + self.addNodes(other.nodes) + +def unityAnimatorToString(nodes): + lines = [] + preamble = """ +%YAML 1.1 +%TAG !u! tag:unity3d.com,2011: +"""[1:][:-1] + lines.append(preamble) + for doc in nodes: + lines.append("--- !u!" + doc.classId() + " &" + doc.anchor) + lines.append(str(doc)) + result = '\n'.join(lines) + + for i in range(0,10): + result = result.replace("\n\n", "\n") + + return result class UnityParser: STREAM_START = 100 @@ -347,137 +317,154 @@ class UnityParser: SEQUENCE_VALUE = 400 - state = STREAM_START - cur_scalar = None - cur_node = None + def __init__(self): + self.state = self.STREAM_START + self.cur_scalar = None + self.cur_node = None - # Simple list of parsed documents. Populated by parse(). - nodes = [] - prev_states = [] + # Simple list of parsed documents. Populated by parse(). + self.nodes = [] + self.prev_states = [] def __str__(self): - lines = [] - preamble = """ -%YAML 1.1 -%TAG !u! tag:unity3d.com,2011: -"""[1:][:-1] - lines.append(preamble) - for doc in self.nodes: - lines.append("--- !u!" + doc.classId() + " &" + doc.anchor) - lines.append(str(doc)) - result = '\n'.join(lines) - - for i in range(0,10): - result = result.replace("\n\n", "\n") - - return result + return unityAnimatorToString(self.nodes) def pushState(self, state): self.prev_states.append(self.state) self.state = state #print("state {} ({})".format(self.state, len(self.prev_states))) - def peekState(self): - return self.state - def popState(self): self.state = self.prev_states[-1] self.prev_states = self.prev_states[0:len(self.prev_states) - 1] #print("state {} ({})".format(self.state, len(self.prev_states))) return self.state + def cleanYaml(self, yaml_str): + lines = [] + first_document = True + for line in yaml_str.split("\n"): + # Add end-of-document indicators. + if line.startswith("---"): + if not first_document: + lines.append("...\n") + first_document = False + + # Remove class ID tag from each block. + if line.startswith("---"): + parts = line.split() + lines.append(parts[0] + " " + parts[2] + "\n") + continue + lines.append(line) + + lines.append("...\n") + return '\n'.join(lines) + def parse(self, yaml_file): + yaml_str = "" with open(yaml_file, "r") as f: yaml_str = f.read() - for event in yaml.parse(yaml_str): - #print("event {}".format(event)) - - if isinstance(event, yaml.StreamStartEvent): - if len(self.prev_states) > 0: - raise Exception("Multiple StreamStartEvents received") - self.pushState(self.STREAM_START) - continue - - if isinstance(event, yaml.StreamEndEvent): - if self.peekState() != self.STREAM_START: - raise Exception("Document end received after state {}".format(self.peekState())) - self.popState() - if len(self.prev_states) > 0: - raise Exception("Extra states at stream end") - continue - - if isinstance(event, yaml.DocumentStartEvent): - if self.peekState() != self.STREAM_START and self.peekState() != self.DOCUMENT_END: - raise Exception("Document start received after state {}".format(self.peekState())) - self.pushState(self.DOCUMENT_START) - continue - - if isinstance(event, yaml.DocumentEndEvent): - if self.peekState() != self.DOCUMENT_START: - raise Exception("Document end received after state {}".format(self.peekState())) - self.popState() - self.nodes.append(self.cur_node) - self.cur_node = None - continue - - if isinstance(event, yaml.MappingStartEvent): - if self.cur_node == None: - self.cur_node = UnityDocument() - self.cur_node.anchor = event.anchor - else: - self.cur_node = self.cur_node.addChildMapping(self.cur_scalar) - self.pushState(self.MAPPING_START) - continue - - if isinstance(event, yaml.MappingEndEvent): - if self.peekState() != self.MAPPING_START: - raise Exception("Mapping end received after state {}".format(self.peekState())) + yaml_str = self.cleanYaml(yaml_str) + + for event in yaml.parse(yaml_str): + if isinstance(event, yaml.StreamStartEvent): + if len(self.prev_states) > 0: + raise Exception("Multiple StreamStartEvents received") + self.pushState(self.STREAM_START) + + elif isinstance(event, yaml.StreamEndEvent): + if self.state != self.STREAM_START: + raise Exception("Document end received after state {}".format(self.state)) + self.popState() + if len(self.prev_states) > 0: + raise Exception("Extra states at stream end") + + elif isinstance(event, yaml.DocumentStartEvent): + if self.state != self.STREAM_START and self.state != self.DOCUMENT_END: + raise Exception("Document start received after state {}".format(self.state)) + self.pushState(self.DOCUMENT_START) + + elif isinstance(event, yaml.DocumentEndEvent): + if self.state != self.DOCUMENT_START: + raise Exception("Document end received after state {}".format(self.state)) + self.popState() + self.nodes.append(self.cur_node) + self.cur_node = None + + elif isinstance(event, yaml.MappingStartEvent): + if self.cur_node == None: + self.cur_node = UnityDocument() + self.cur_node.anchor = event.anchor + else: + self.cur_node = self.cur_node.addChildMapping(self.cur_scalar) + self.pushState(self.MAPPING_START) + + elif isinstance(event, yaml.MappingEndEvent): + if self.state != self.MAPPING_START: + raise Exception("Mapping end received after state {}".format(self.state)) + self.popState() + if self.state == self.MAPPING_KEY: self.popState() - if self.peekState() == self.MAPPING_KEY: - self.popState() - if self.cur_node.parent != None: - self.cur_node = self.cur_node.parent - continue - - if isinstance(event, yaml.SequenceStartEvent): - self.cur_node = self.cur_node.addChildSequence(self.cur_scalar) - self.pushState(self.SEQUENCE_VALUE) - continue - - if isinstance(event, yaml.SequenceEndEvent): - if self.peekState() != self.SEQUENCE_VALUE: - raise Exception("Sequence end received after state {}".format(self.peekState())) - self.popState() - if self.peekState() == self.MAPPING_KEY: - self.popState() + if self.cur_node.parent != None: self.cur_node = self.cur_node.parent - continue - - if isinstance(event, yaml.ScalarEvent): - if self.peekState() == self.MAPPING_START: - self.cur_scalar = event.value - self.pushState(self.MAPPING_KEY) - elif self.peekState() == self.MAPPING_KEY: - self.cur_node.mapping[self.cur_scalar] = event.value - self.popState() - elif self.peekState() == self.SEQUENCE_VALUE: - self.cur_node.sequence.append(event.value) - else: - raise Exception("Scalar event received after state {}".format(self.peekState())) - continue + elif isinstance(event, yaml.SequenceStartEvent): + self.cur_node = self.cur_node.addChildSequence(self.cur_scalar) + self.pushState(self.SEQUENCE_VALUE) + + elif isinstance(event, yaml.SequenceEndEvent): + if self.state != self.SEQUENCE_VALUE: + raise Exception("Sequence end received after state {}".format(self.state)) + self.popState() + if self.state == self.MAPPING_KEY: + self.popState() + self.cur_node = self.cur_node.parent + + elif isinstance(event, yaml.ScalarEvent): + if self.state == self.MAPPING_START: + self.cur_scalar = event.value + self.pushState(self.MAPPING_KEY) + elif self.state == self.MAPPING_KEY: + self.cur_node.mapping[self.cur_scalar] = event.value + self.popState() + elif self.state == self.SEQUENCE_VALUE: + self.cur_node.sequence.append(event.value) + else: + raise Exception("Scalar event received after state {}".format(self.state)) + else: raise Exception("Unhandled event {}".format(event)) + continue if __name__ == "__main__": - arg0 = "tst.fx" - tmp0 = "fx0.controller" - importUnityYaml(arg0, tmp0) + arg0 = "../FX.controller" + print("Parsing {}".format(arg0), file=sys.stderr) + parser0 = UnityParser() + try: + parser0.parse(arg0) + except Exception as e: + print("exception: {}".format(e)) + + anim0 = UnityAnimator() + anim0.addNodes(parser0.nodes) - parser = UnityParser() + arg1 = "TaSTT_fx.controller" + print("Parsing {}".format(arg1), file=sys.stderr) + parser1 = UnityParser() try: - parser.parse(tmp0) + parser1.parse(arg1) except Exception as e: print("exception: {}".format(e)) - print(parser) + + anim1 = UnityAnimator() + anim1.addNodes(parser1.nodes) + + #anim1.nodes = [] + #anim1.id_to_node = {} + print("Merging animators", file=sys.stderr) + anim0.merge(anim1) + + #print(unityAnimatorToString(parser0.nodes)) + print("Serializing", file=sys.stderr) + print(unityAnimatorToString(anim0.nodes)) |