Add simple yaml parser (WIP)

Add parser for Unity's malformed YAML. This should make it easier to
manipulate animators.

It probably doesn't quite work yet, and certainly needs some usability
features.

1 files changed, 326 insertions, 3 deletions

diff --git a/libunity.py b/libunity.py
index 3470ea0..5559c5f 100644
--- a/libunity.py
+++ b/libunity.py
@@ -1,5 +1,9 @@
 #!/usr/bin/env python3
 
+import enum
+# python3 -m pip install pyyaml
+import yaml
+
 def getObjectIds(fx):
     fx_ids = set()
     with open(fx, "r") as f:
@@ -153,8 +157,327 @@ def mergeFX(fx0, fx1):
 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
+        # mappings have an anchor.
+        self.anchor = None
+
+        # Pointer to the Node containing this one.
+        self.parent = None
+
+class Sequence(Node):
+    def __init__(self):
+        super().__init__()
+        self.sequence = []
+
+    def prettyPrint(self, first_indent=None, leading_newline=None):
+        depth = 0
+        p = self.parent
+        while p != None:
+            depth += 1
+            p = p.parent
+        indent = "  " * depth
+
+        lines = []
+        first = True
+        for item in self.sequence:
+            cur_indent = indent
+            if first:
+                if first_indent != None:
+                    cur_indent = first_indent
+                first = False
+            if hasattr(item, "prettyPrint"):
+                lines.append("{}- {}".format(cur_indent, item.prettyPrint(first_indent="", leading_newline=False)))
+            else:
+                lines.append("{}- {}".format(cur_indent, item))
+
+        if len(lines) == 0:
+            return "[]"
+
+        return "\n" + '\n'.join(lines)
+
+    def __str__(self):
+        return self.prettyPrint()
+
+    def addChildMapping(self, anchor = None):
+        child = Mapping()
+        child.anchor = anchor
+        child.parent = self
+        child.sequence = []
+
+        self.sequence.append(child)
+
+        return child
+
+    def addChildSequence(self, anchor = None):
+        child = Sequence()
+        child.anchor = anchor
+        child.parent = self
+        child.sequence = []
+
+        self.sequence.append(child)
+
+        return child
+
+class Mapping(Node):
+    def __init__(self):
+        super().__init__()
+        self.mapping = {}
+
+    def prettyPrint(self, first_indent=None, leading_newline=True):
+        depth = 0
+        p = self.parent
+        while p != None:
+            depth += 1
+            p = p.parent
+        indent = "  " * depth
+
+        lines = []
+        first = True
+        for k, v in self.mapping.items():
+            cur_indent = indent
+            if first:
+                if first_indent != None:
+                    cur_indent = first_indent
+                first = False
+            lines.append("{}{}: {}".format(cur_indent, k, v))
+
+        result = '\n'.join(lines)
+
+        # Inline 1-item mappings, matching Unity behavior.
+        if len(self.mapping.keys()) == 1 and len(result.split("\n")) == 1:
+            if first_indent == None:
+                return self.prettyPrint(first_indent="")
+            return "{" + lines[0] + "}"
+
+        if leading_newline:
+            result = "\n" + result
+
+        return result
+
+    def __str__(self):
+        return self.prettyPrint()
+
+    def addChildMapping(self, key, anchor = None):
+        child = Mapping()
+        child.anchor = anchor
+        child.parent = self
+        child.mapping = {}
+
+        self.mapping[key] = child
+
+        return child
+
+    def addChildSequence(self, key, anchor = None):
+        child = Sequence()
+        child.anchor = anchor
+        child.parent = self
+        child.mapping = {}
+
+        self.mapping[key] = child
+
+        return child
+
+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))
+
+class UnityParser:
+    STREAM_START = 100
+    STREAM_END = 199
+
+    DOCUMENT_START = 200
+    DOCUMENT_END = 299
+
+    MAPPING_START = 300
+    MAPPING_KEY = 301
+
+    SEQUENCE_VALUE = 400
+
+    state = STREAM_START
+    cur_scalar = None
+    cur_node = None
+
+    # Simple list of parsed documents. Populated by parse().
+    nodes = []
+    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
+
+    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 parse(self, yaml_file):
+        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()))
+                    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()
+                    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
+
+                raise Exception("Unhandled event {}".format(event))
+
 if __name__ == "__main__":
-    #mergeFX("FX.controller", "TaSTT/generated_fx.controller")
-    #print(extractFirstBlock("../FX.controller", "--- !u!91", "---"))
-    print(mergeFX("TaSTT_fx.controller", "../FX.controller"))
+    arg0 = "tst.fx"
+
+    tmp0 = "fx0.controller"
+    importUnityYaml(arg0, tmp0)
+
+    parser = UnityParser()
+    try:
+        parser.parse(tmp0)
+    except Exception as e:
+        print("exception: {}".format(e))
+    print(parser)