Fix --tyro-write-completion when use_underscores=True (#173)

* Fix `--tyro-write-completion` when `use_underscores=True`

* Revert example update

src/tyro/_cli.py

@@ -381,8 +381,14 @@ def _cli_impl(
     #
     # Note: --tyro-print-completion is deprecated! --tyro-write-completion is less prone
     # to errors from accidental logging, print statements, etc.
-    print_completion = len(args) >= 2 and args[0] == "--tyro-print-completion"
-    write_completion = len(args) >= 3 and args[0] == "--tyro-write-completion"
+    print_completion = False
+    write_completion = False
+    if len(args) >= 2:
+        # We replace underscores with hyphens to accomodate for `use_undercores`.
+        print_completion = args[0].replace("_", "-") == "--tyro-print-completion"
+        write_completion = (
+            len(args) >= 3 and args[0].replace("_", "-") == "--tyro-write-completion"
+        )
 
     # Note: setting USE_RICH must happen before the parser specification is generated.
     # TODO: revisit this. Ideally we should be able to eliminate the global state
@@ -442,7 +448,7 @@ def _cli_impl(
                 f" {completion_shell}"
             )
 
-            if write_completion:
+            if write_completion and completion_target_path != pathlib.Path("-"):
                 assert completion_target_path is not None
                 completion_target_path.write_text(
                     shtab.complete(

tests/helptext_utils.py

@@ -44,21 +44,35 @@ def get_helptext_with_checks(
     unformatted_usage = parser.format_usage()
     assert tyro._strings.strip_ansi_sequences(unformatted_usage) == unformatted_usage
 
-    # Completion scripts; just smoke test for now.
-    with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")):
-        # `--tyro-print-completion` is deprecated! We should use `--tyro-write-completion` instead.
-        tyro.cli(f, default=default, args=["--tyro-print-completion", "bash"])
-    with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")):
-        # `--tyro-print-completion` is deprecated! We should use `--tyro-write-completion` instead.
-        tyro.cli(f, default=default, args=["--tyro-print-completion", "zsh"])
-    with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")):
+    # Basic checks for completion scripts.
+    with pytest.raises(SystemExit):
         tyro.cli(
             f, default=default, args=["--tyro-write-completion", "bash", os.devnull]
         )
-    with pytest.raises(SystemExit), contextlib.redirect_stdout(open(os.devnull, "w")):
-        tyro.cli(
-            f, default=default, args=["--tyro-write-completion", "zsh", os.devnull]
-        )
+    for shell in ["bash", "zsh"]:
+        for command in ["--tyro-print-completion", "--tyro-write-completion"]:
+            target = io.StringIO()
+            with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
+                if command == "--tyro-write-completion":
+                    tyro.cli(f, default=default, args=[command, shell, "-"])
+                else:
+                    # `--tyro-print-completion` is deprecated! We should use `--tyro-write-completion` instead.
+                    tyro.cli(f, default=default, args=[command, shell])
+            output = target.getvalue()
+            assert "shtab" in output
+
+    # Test with underscores
+    for shell in ["bash", "zsh"]:
+        target = io.StringIO()
+        with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
+            tyro.cli(
+                f,
+                default=default,
+                args=["--tyro_write_completion", shell, "-"],
+                use_underscores=True,
+            )
+        output = target.getvalue()
+        assert "shtab" in output
 
     # Get the actual helptext.
     target = io.StringIO()