More functionality in dump

trustgraph-cli/scripts/tg-dump-msgpack

@@ -10,7 +10,7 @@ import msgpack
 import sys
 import argparse
 
-def run(input_file):
+def dump(input_file, action):
 
     with open(input_file, 'rb') as f:
 
@@ -19,6 +19,45 @@ def run(input_file):
         for unpacked in unpacker:
             print(unpacked)
 
+def summary(input_file, action):
+
+    vector_dim = None
+
+    triples = set()
+
+    max_records = 10
+
+    with open(input_file, 'rb') as f:
+
+        unpacker = msgpack.Unpacker(f, raw=False)
+
+        rec_count = 0
+
+        for msg in unpacker:
+
+            if msg[0] == "ge":
+                vector_dim = len(msg[1]["v"][0])
+
+            if msg[0] == "t":
+
+                for elt in msg[1]["m"]["m"]:
+                    triples.add(
+                        elt["s"]["v"],
+                        elt["p"]["v"],
+                        elt["o"]["v"],
+                    )
+
+            if rec_count > max_records: break
+            rec_count += 1
+
+    print("Vector dimension:", vector_dim)
+    print(triples)
+
+#     print(f"Graph-embeddings:   {ge:10d}")
+#     for k, v in sorted(ge_dim.items()):
+#         print(f"    Dimension {k}:  {v:10d}")
+#     print(f"Triples:            {t:10d}")
+
 def main():
     
     parser = argparse.ArgumentParser(
@@ -32,9 +71,24 @@ def main():
         help=f'Input file'
     )
 
+    parser.add_argument(
+        '-s', '--summary', action="store_const", const="summary",
+        dest="action",
+        help=f'Show a summary'
+    )
+
+    parser.add_argument(
+        '-r', '--records', action="store_const", const="records",
+        dest="action",
+        help=f'Dump individual records'
+    )
+
     args = parser.parse_args()
 
-    run(**vars(args))
+    if args.action == "summary":
+        summary(**vars(args))
+    else:
+        dump(**vars(args))
 
 main()