Repository: omni-us/jsonargparse

Issue:
Dynamic selection of function given an argument fails with an `AttributeError`

## 🐛 Bug report

I am trying to setup a CLI system where I have a multitude of subcommands (only 1 in the repro code, `finetune`) and some of them take a `--method` argument that allows selecting a different function.

See https://github.com/Lightning-AI/litgpt/issues/996 for the real purpose of this.

### To reproduce

```python
from typing import Literal


def methodA(foo: int = 111):
    print(locals())


def methodB(bar: int = 123):
    print(locals())


def finetune(method: Literal["A", "B"], *args, **kwargs):
    if method == "A":
        return methodA(*args, **kwargs)
    elif method == "B":
        return methodB(*args, **kwargs)
    raise NotImplementedError(method)


if __name__ == "__main__":
    from jsonargparse import CLI

    CLI(finetune)
```

```python
Traceback (most recent call last):
  File "/home/carmocca/git/repro.py", line 19, in <module>
    CLI(finetune)
  File "/home/carmocca/git/nightly-venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 88, in CLI
    _add_component_to_parser(components, parser, as_positional, fail_untyped, config_help)
  File "/home/carmocca/git/nightly-venv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 184, in _add_component_to_parser
    added_args = parser.add_function_arguments(component, as_group=False, **kwargs)
  File "/home/carmocca/git/nightly-venv/lib/python3.10/site-packages/jsonargparse/_signatures.py", line 190, in add_function_arguments
    return self._add_signature_arguments(
  File "/home/carmocca/git/nightly-venv/lib/python3.10/site-packages/jsonargparse/_signatures.py", line 270, in _add_signature_arguments
    self._add_signature_parameter(
  File "/home/carmocca/git/nightly-venv/lib/python3.10/site-packages/jsonargparse/_signatures.py", line 343, in _add_signature_parameter
    if group_name in group.parser.groups:
AttributeError: 'ArgumentParser' object has no attribute 'parser'
```

### Expected behavior

The ideal behaviour would be that `python repro.py -h` only lists the method argument:

```python
positional arguments:
  {A,B}                 (required, type: Literal['A', 'B'])
```

And then when one does `python repro.py --method B -h` then it also shows:

```python
options:
  -h, --help            Show this help message and exit.
  --config CONFIG       Path to a configuration file.
  --print_config[=flags]
                        Print the configuration after applying all other arguments and exit. The optional flags customizes the output and are one or more keywords separated by comma. The supported flags are: comments, skip_default,
                        skip_null.
  --bar BAR             (type: int, default: 123)
```

I don't know if the AST is smart enough to parse the `*args, **kwargs` specific to the method selected.

I'm happy to try out a different way to support CLI behaviour like this, if there are better ways to approach this problem.

### Environment

- jsonargparse version (e.g., 4.8.0): 4.27.5


Comments:

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index b6bc59b..99d9f6a 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -26,6 +26,8 @@ Fixed
   produces an invalid string default.
 - dataclass single parameter change incorrectly resetting previous values (`#464
   <https://github.com/omni-us/jsonargparse/issues/464>`__).
+- Add function signature failing when conditionally calling different functions
+  (`#467 <https://github.com/omni-us/jsonargparse/issues/467>`__).
 
 
 v4.27.5 (2024-02-12)
diff --git a/jsonargparse/_signatures.py b/jsonargparse/_signatures.py
index 0f83730..94ca7af 100644
--- a/jsonargparse/_signatures.py
+++ b/jsonargparse/_signatures.py
@@ -3,7 +3,7 @@
 import dataclasses
 import inspect
 import re
-from argparse import SUPPRESS
+from argparse import SUPPRESS, ArgumentParser
 from contextlib import suppress
 from typing import Any, Callable, List, Optional, Set, Tuple, Type, Union
 
@@ -255,7 +255,7 @@ class SignatureArguments(LoggerProperty):
         ## Create group if requested ##
         doc_group = get_doc_short_description(function_or_class, method_name, logger=self.logger)
         component = getattr(function_or_class, method_name) if method_name else function_or_class
-        group = self._create_group_if_requested(
+        container = self._create_group_if_requested(
             component,
             nested_key,
             as_group,
@@ -268,7 +268,7 @@ class SignatureArguments(LoggerProperty):
         added_args: List[str] = []
         for param in params:
             self._add_signature_parameter(
-                group,
+                container,
                 nested_key,
                 param,
                 added_args,
@@ -283,7 +283,7 @@ class SignatureArguments(LoggerProperty):
 
     def _add_signature_parameter(
         self,
-        group,
+        container,
         nested_key: Optional[str],
         param,
         added_args: List[str],
@@ -339,11 +339,14 @@ class SignatureArguments(LoggerProperty):
         dest = (nested_key + "." if nested_key else "") + name
         args = [dest if is_required and as_positional else "--" + dest]
         if param.origin:
+            parser = container
+            if not isinstance(container, ArgumentParser):
+                parser = getattr(container, "parser")
             group_name = "; ".join(str(o) for o in param.origin)
-            if group_name in group.parser.groups:
-                group = group.parser.groups[group_name]
+            if group_name in parser.groups:
+                container = parser.groups[group_name]
             else:
-                group = group.parser.add_argument_group(
+                container = parser.add_argument_group(
                     f"Conditional arguments [origins: {group_name}]",
                     name=group_name,
                 )
@@ -372,7 +375,7 @@ class SignatureArguments(LoggerProperty):
                     args=args,
                     kwargs=kwargs,
                     enable_path=enable_path,
-                    container=group,
+                    container=container,
                     logger=self.logger,
                     sub_add_kwargs=sub_add_kwargs,
                 )
@@ -387,7 +390,7 @@ class SignatureArguments(LoggerProperty):
             if is_dataclass_like_typehint:
                 kwargs.update(sub_add_kwargs)
             with ActionTypeHint.allow_default_instance_context():
-                action = group.add_argument(*args, **kwargs)
+                action = container.add_argument(*args, **kwargs)
             action.sub_add_kwargs = sub_add_kwargs
             if is_subclass_typehint and len(subclass_skip) > 0:
                 action.sub_add_kwargs["skip"] = subclass_skip


Test patch:
diff --git a/jsonargparse_tests/test_cli.py b/jsonargparse_tests/test_cli.py
index 2590170..0bb933b 100644
--- a/jsonargparse_tests/test_cli.py
+++ b/jsonargparse_tests/test_cli.py
@@ -14,6 +14,7 @@ import yaml
 
 from jsonargparse import CLI, capture_parser, lazy_instance
 from jsonargparse._optionals import docstring_parser_support, ruyaml_support
+from jsonargparse._typehints import Literal
 from jsonargparse.typing import final
 from jsonargparse_tests.conftest import skip_if_docstring_parser_unavailable
 
@@ -120,6 +121,31 @@ def test_multiple_functions_subcommand_help():
     assert "--a2 A2" in out
 
 
+def conditionalA(foo: int = 1):
+    return foo
+
+
+def conditionalB(bar: int = 2):
+    return bar
+
+
+def conditional_function(fn: "Literal['A', 'B']", *args, **kwargs):
+    if fn == "A":
+        return conditionalA(*args, **kwargs)
+    elif fn == "B":
+        return conditionalB(*args, **kwargs)
+    raise NotImplementedError(fn)
+
+
+@pytest.mark.skipif(condition=sys.version_info < (3, 9), reason="python>=3.9 is required")
+@pytest.mark.skipif(condition=not Literal, reason="Literal is required")
+def test_literal_conditional_function():
+    out = get_cli_stdout(conditional_function, args=["--help"])
+    assert "Conditional arguments" in out
+    assert "--foo FOO             (type: int, default: Conditional<ast-resolver> {1, NOT_ACCEPTED})" in out
+    assert "--bar BAR             (type: int, default: Conditional<ast-resolver> {2, NOT_ACCEPTED})" in out
+
+
 # single class tests