Add libunity CLI

* Add guid scanning method
  * Generate mapping from guid to filename
  * Mapping may be saved & restored from disk

1 files changed, 83 insertions, 31 deletions

diff --git a/libunity.py b/libunity.py
index 097e8b8..830cc08 100644
--- a/libunity.py
+++ b/libunity.py
@@ -1,7 +1,10 @@
 #!/usr/bin/env python3
 
+import argparse
 import copy
 import enum
+import os
+import pickle
 import sys
 # python3 -m pip install pyyaml
 import yaml
@@ -435,36 +438,85 @@ class UnityParser:
                 raise Exception("Unhandled event {}".format(event))
             continue
 
-if __name__ == "__main__":
+def getGuidMap(d):
+    result = {}
+    for f in os.scandir(d):
+        path = f.path
+        if f.is_dir():
+            result.update(getGuidMap(path))
+        if not f.is_file():
+            continue
+        suffix = ".meta"
+        if path.endswith(suffix):
+            with open(path, "r") as f:
+                for line in f:
+                    if line.startswith("guid"):
+                        guid = line.split()[1]
+                        result[path[:-len(suffix)]] = guid
+    return result
 
-    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)
-
-    arg1 = "TaSTT_fx.controller"
-    print("Parsing {}".format(arg1), file=sys.stderr)
-    parser1 = UnityParser()
-    try:
-        parser1.parse(arg1)
-    except Exception as e:
-        print("exception: {}".format(e))
-
-    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))
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("cmd", type=str)
+    parser.add_argument("--fx0", type=str, help="The first animator to merge")
+    parser.add_argument("--fx1", type=str, help="The second animator to merge")
+    parser.add_argument("--project_root", type=str, help="The path to the " +
+            "Unity project Assets folder")
+    parser.add_argument("--save_to", type=str, help="The path to save the " +
+            "result of the computation")
+    parser.add_argument("--guid_map", type=str, help="Path to guid.map, " +
+            "generated by a previous call to `guid_map`")
+    args = parser.parse_args()
+
+    guid_map = {}
+    if args.guid_map:
+        with open(args.guid_map, 'rb') as f:
+            guid_map = pickle.load(f)
+
+    if args.cmd == "merge":
+        if not args.fx0 or not args.fx1:
+            print("Usage: ./libunity.py merge --fx0 my_fx.controller --fx1 " +
+                    "my_other_fx.controller")
+            sys.exit(1)
+
+        print("Parsing {}".format(args.fx0), file=sys.stderr)
+        parser0 = UnityParser()
+        try:
+            parser0.parse(args.fx0)
+        except Exception as e:
+            print("exception: {}".format(e))
+
+        anim0 = UnityAnimator()
+        anim0.addNodes(parser0.nodes)
+
+        arg1 = "TaSTT_fx.controller"
+        print("Parsing {}".format(args.fx1), file=sys.stderr)
+        parser1 = UnityParser()
+        try:
+            parser1.parse(args.fx1)
+        except Exception as e:
+            print("exception: {}".format(e))
+
+        anim1 = UnityAnimator()
+        anim1.addNodes(parser1.nodes)
+
+        print("Merging animators", file=sys.stderr)
+        anim0.merge(anim1)
+
+        print("Serializing", file=sys.stderr)
+        print(unityAnimatorToString(anim0.nodes))
+    elif args.cmd == "guid_map":
+        if not args.project_root or not args.save_to:
+            print("Usage: ./libunity.py guid_map --project_root " +
+                    "/path/to/unity/assets --save_to /path/to/guid.map")
+            sys.exit(1)
+
+        print("Looking up GUIDs under {}".format(args.project_root),
+                file=sys.stderr)
+        guid_map = getGuidMap(args.project_root)
+        print("Saving to {}".format(args.save_to), file=sys.stderr)
+        with open(args.save_to, 'wb') as f:
+            pickle.dump(guid_map, f)
+    else:
+        print("Unrecognized command: {}".format(args.cmd))