.

wcfs/testprog/treegen.py

@@ -315,6 +315,18 @@ def cmd_allstructs(argv):
     AllStructs(zstor, kv1, kv2, n)
 
 
+@func
+def cmd_trees(argv):
+    if len(argv) != 1:
+        print("Usage: cat trees |treegen trees <zurl>", file=sys.stderr)
+        sys.exit(1)
+
+    zurl  = argv[0]
+    zstor = storageFromURL(zurl)
+    defer(zstor.close)
+    Trees(zstor, sys.stdin)
+
+
 cmdRegistry = {
     'allstructs':   cmd_allstructs,
     'trees':        cmd_trees,
@@ -325,7 +337,10 @@ def main():
         print("Usage: treegen <command> ...", file=sys.stderr)
         sys.exit(1)
 
-    cmd  = cmdRegistry[sys.argv[1]]
+    cmd  = cmdRegistry.get(sys.argv[1], None)
+    if cmd is None:
+        print("E: treegen: unknown command %r" % sys.argv[1])
+        sys.exit(1)
     argv = sys.argv[2:]
     cmd(argv)