From 1af508002e4117cff12e90e8a3c11cb219daa31b Mon Sep 17 00:00:00 2001
From: Jeremy Leibs <jeremy@rerun.io>
Date: Wed, 19 Apr 2023 16:12:29 -0400
Subject: [PATCH] Add argument parsing to the rerun_demo (#1925)

---
 rerun_py/rerun_sdk/rerun_demo/__main__.py | 34 ++++++++++++++++++-----
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/rerun_py/rerun_sdk/rerun_demo/__main__.py b/rerun_py/rerun_sdk/rerun_demo/__main__.py
index ce5a5b34fdf4..dcb58f23d1ef 100644
--- a/rerun_py/rerun_sdk/rerun_demo/__main__.py
+++ b/rerun_py/rerun_sdk/rerun_demo/__main__.py
@@ -5,15 +5,16 @@
 import sys
 
 
-def run_cube():
+def run_cube(args: argparse.Namespace):
     import math
 
     import numpy as np
     import rerun as rr
 
-    rr.init("Cube", spawn=True, default_enabled=True)
     from rerun_demo.data import build_color_grid
 
+    rr.script_setup(args, "Cube")
+
     STEPS = 100
     twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
     for t in range(STEPS):
@@ -21,10 +22,25 @@ def run_cube():
         cube = build_color_grid(10, 10, 10, twist=twists[t])
         rr.log_points("cube", positions=cube.positions, colors=cube.colors, radii=0.5)
 
+    rr.script_teardown(args)
+
 
-def run_colmap():
+def run_colmap(args):
     from rerun import bindings, unregister_shutdown  # type: ignore[attr-defined]
 
+    serve_opts = []
+
+    # TODO(https://github.com/rerun-io/rerun/issues/1924): The need to special-case
+    # this flag conversion is a bit awkward.
+    if args.connect or args.addr:
+        print("Connecting to external viewer is only supported with the --cube demo.", file=sys.stderr)
+        exit(1)
+    if args.save:
+        print("Saving an RRD file is only supported from the --cube demo.", file=sys.stderr)
+        exit(1)
+    if args.serve:
+        serve_opts.append("--web-viewer")
+
     # We don't need to call shutdown in this case. Rust should be handling everything
     unregister_shutdown()
 
@@ -33,11 +49,13 @@ def run_colmap():
         print("No demo file found at {}. Package was built without demo support".format(rrd_file), file=sys.stderr)
         exit(1)
     else:
-        exit(bindings.main([sys.argv[0], str(rrd_file)]))
+        exit(bindings.main([sys.argv[0], str(rrd_file)] + serve_opts))
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(description="Run rerun example programs")
+    import rerun as rr
+
+    parser = argparse.ArgumentParser(description="Run rerun example programs.")
 
     group = parser.add_mutually_exclusive_group()
 
@@ -53,16 +71,18 @@ def main() -> None:
         help="Run the COLMAP data demo",
     )
 
+    rr.script_add_args(parser)
+
     args = parser.parse_args()
 
     if not any([args.cube, args.colmap]):
         args.cube = True
 
     if args.cube:
-        run_cube()
+        run_cube(args)
 
     elif args.colmap:
-        run_colmap()
+        run_colmap(args)
 
 
 if __name__ == "__main__":