From 3b7e37ddbb3a261db9c298cc378cc7f919f3755d Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Thu, 10 Oct 2024 14:18:59 -0700
Subject: [PATCH] Fix `--tyro-write-completion` when `use_underscores=True`
 (#173)

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

* Revert example update
---
 src/tyro/_cli.py        | 12 +++++++++---
 tests/helptext_utils.py | 38 ++++++++++++++++++++++++++------------
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/src/tyro/_cli.py b/src/tyro/_cli.py
index 4173b16f..ff01894c 100644
--- a/src/tyro/_cli.py
+++ b/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(
diff --git a/tests/helptext_utils.py b/tests/helptext_utils.py
index 041f5c58..73dd1fc9 100644
--- a/tests/helptext_utils.py
+++ b/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()