Improve prefix omitting flag keep groups in help (#252)

* Improve prefix omitting flag keep groups in help

* ruff

---------

Co-authored-by: brentyi <yibrenth@gmail.com>

Author: e3rd

src/tyro/_parsers.py

@@ -247,16 +247,6 @@ def apply_args(
         def format_group_name(group_name: str) -> str:
             return (group_name + " options").strip()
 
-        def group_name_from_arg(arg: _arguments.ArgumentDefinition) -> str:
-            prefix = arg.lowered.name_or_flags[0]
-            if prefix.startswith("--"):
-                prefix = prefix[2:]
-            if "." in prefix:
-                prefix = prefix.rpartition(".")[0]
-            else:
-                prefix = ""
-            return prefix
-
         group_from_group_name: Dict[str, argparse._ArgumentGroup] = {
             "": parser._action_groups[1],
             **{
@@ -274,7 +264,7 @@ def group_name_from_arg(arg: _arguments.ArgumentDefinition) -> str:
             if arg.is_suppressed():
                 continue
 
-            group_name = group_name_from_arg(arg)
+            group_name = arg.extern_prefix
             if group_name not in group_from_group_name:
                 description = (
                     parent.helptext_from_intern_prefixed_field_name.get(

tests/test_conf.py

@@ -1112,9 +1112,15 @@ class TrainConfig:
         args="--model.num-slots 3".split(" "),
     ) == TrainConfig(ModelConfig(num_slots=3))
 
-    assert tyro.cli(
-        tyro.conf.OmitArgPrefixes[TrainConfig], args="--num-slots 3".split(" ")
-    ) == TrainConfig(ModelConfig(num_slots=3))
+    annot = tyro.conf.OmitArgPrefixes[TrainConfig]
+    assert tyro.cli(annot, args="--num-slots 3".split(" ")) == TrainConfig(
+        ModelConfig(num_slots=3)
+    )
+
+    # groups are still printed in the help text
+    help_text = get_helptext_with_checks(annot)
+    assert "model options" in help_text
+    assert "--num-slots" in help_text
 
 
 def test_custom_constructor_0() -> None:
@@ -1539,8 +1545,8 @@ class DatasetConfig:
     with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
         instantiate_dataclasses((OptimizerConfig, DatasetConfig), args=["--help"])
     helptext = target.getvalue()
-    assert "OptimizerConfig options" not in helptext
-    assert "DatasetConfig options" not in helptext
+    assert "OptimizerConfig options" in helptext
+    assert "DatasetConfig options" in helptext
 
 
 def test_counter_action() -> None:

tests/test_py311_generated/test_conf_generated.py

@@ -1546,8 +1546,8 @@ class DatasetConfig:
     with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
         instantiate_dataclasses((OptimizerConfig, DatasetConfig), args=["--help"])
     helptext = target.getvalue()
-    assert "OptimizerConfig options" not in helptext
-    assert "DatasetConfig options" not in helptext
+    assert "OptimizerConfig options" in helptext
+    assert "DatasetConfig options" in helptext
 
 
 def test_counter_action() -> None: