Fix: List of union of classes not accepted by add_subclass_arguments in python>=3.11 (#522)

Author: mauvilsa

CHANGELOG.rst

@@ -34,6 +34,9 @@ Fixed
   (`#523 <https://github.com/omni-us/jsonargparse/pull/523>`__).
 - Failing to parse list of dataclasses with nested optional dataclass (`#527
   <https://github.com/omni-us/jsonargparse/pull/527>`__).
+- List of union of classes not accepted by ``add_subclass_arguments`` in
+  ``python>=3.11`` (`#522
+  <https://github.com/omni-us/jsonargparse/pull/522>`__).
 
 
 v4.29.0 (2024-05-24)

jsonargparse/_actions.py

@@ -349,14 +349,9 @@ def __init__(self, baseclass=None, **kwargs):
             super().__init__(**kwargs)
 
     def update_init_kwargs(self, kwargs):
-        if get_typehint_origin(self._baseclass) == Union:
-            from ._typehints import ActionTypeHint
+        from ._typehints import get_subclasses_from_type
 
-            self._basename = iter_to_set_str(
-                c.__name__ for c in self._baseclass.__args__ if ActionTypeHint.is_subclass_typehint(c)
-            )
-        else:
-            self._basename = self._baseclass.__name__
+        self._basename = iter_to_set_str(get_subclasses_from_type(self._baseclass))
         kwargs.update(
             {
                 "metavar": "CLASS_PATH_OR_NAME",

jsonargparse/_parameter_resolvers.py

@@ -28,6 +28,7 @@
 from ._util import (
     ClassFromFunctionBase,
     get_import_path,
+    get_typehint_origin,
     iter_to_set_str,
     unique,
 )
@@ -74,6 +75,10 @@ def __init__(self, resolver: str, data: Any) -> None:
 
 
 def get_parameter_origins(component, parent) -> Optional[str]:
+    from ._typehints import get_subclasses_from_type, sequence_origin_types
+
+    if get_typehint_origin(component) in sequence_origin_types:
+        component = get_subclasses_from_type(component, names=False)
     if isinstance(component, tuple):
         assert parent is None or len(component) == len(parent)
         return iter_to_set_str(get_parameter_origins(c, parent[n] if parent else None) for n, c in enumerate(component))

jsonargparse/_signatures.py

@@ -26,9 +26,10 @@
     ActionTypeHint,
     LazyInitBaseClass,
     callable_instances,
+    get_subclasses_from_type,
     is_optional,
 )
-from ._util import get_import_path, get_private_kwargs, iter_to_set_str
+from ._util import get_private_kwargs, iter_to_set_str
 from .typing import register_pydantic_type
 
 __all__ = [
@@ -510,8 +511,8 @@ def add_subclass_arguments(
             raise ValueError("Not allowed for dataclass-like classes.")
         if type(baseclass) is not tuple:
             baseclass = (baseclass,)  # type: ignore[assignment]
-        if not baseclass or not all(inspect.isclass(c) for c in baseclass):
-            raise ValueError(f"Expected 'baseclass' argument to be a class or a tuple of classes: {baseclass}")
+        if not baseclass or not all(ActionTypeHint.is_subclass_typehint(c, also_lists=True) for c in baseclass):
+            raise ValueError(f"Expected 'baseclass' to be a subclass type or a tuple of subclass types: {baseclass}")
 
         doc_group = None
         if len(baseclass) == 1:  # type: ignore[arg-type]
@@ -530,7 +531,7 @@ def add_subclass_arguments(
         if skip is not None:
             skip = {f"{nested_key}.init_args." + s for s in skip}
         param = ParamData(name=nested_key, annotation=Union[baseclass], component=baseclass)
-        str_baseclass = iter_to_set_str(get_import_path(x) for x in baseclass)
+        str_baseclass = iter_to_set_str(get_subclasses_from_type(param.annotation))
         kwargs.update(
             {
                 "metavar": metavar,

jsonargparse/_typehints.py

@@ -668,6 +668,22 @@ def is_pathlike(typehint) -> bool:
     return is_subclass(typehint, os.PathLike)
 
 
+def get_subclasses_from_type(typehint, names=True, subclasses=None) -> tuple:
+    if subclasses is None:
+        subclasses = []
+    origin = get_typehint_origin(typehint)
+    if origin == Union or origin in sequence_origin_types:
+        for subtype in typehint.__args__:
+            get_subclasses_from_type(subtype, names, subclasses)
+    elif ActionTypeHint.is_subclass_typehint(typehint, all_subtypes=False):
+        if names:
+            if typehint.__name__ not in subclasses:
+                subclasses.append(typehint.__name__)
+        elif typehint not in subclasses:
+            subclasses.append(typehint)
+    return tuple(subclasses)
+
+
 def raise_unexpected_value(message: str, val: Any = inspect._empty, exception: Optional[Exception] = None) -> NoReturn:
     if val is not inspect._empty:
         message += f". Got value: {val}"
@@ -1105,14 +1121,20 @@ def is_private(class_path):
         return "._" in class_path
 
     def add_subclasses(cl):
+        if hasattr(cl, "__args__") and get_typehint_origin(cl) in {List, list, Union}:
+            for arg in cl.__args__:
+                add_subclasses(arg)
+            return
         try:
             class_path = get_import_path(cl)
         except (ImportError, AttributeError) as err:  # Attribute is added in case of dot notation imports
             warning(f"Hit failing import with following error: {err}")
             return
-        if is_local(cl) or issubclass(cl, LazyInitBaseClass):
+        if is_local(cl) or is_subclass(cl, LazyInitBaseClass):
             return
         if not (inspect.isabstract(cl) or is_private(class_path)):
+            if class_path in subclass_list:
+                return
             subclass_list.append(class_path)
         for subclass in cl.__subclasses__() if hasattr(cl, "__subclasses__") else []:
             add_subclasses(subclass)
@@ -1124,7 +1146,7 @@ def add_subclasses(cl):
 
     if get_typehint_origin(cls) in {Union, Type, type}:
         for arg in cls.__args__:
-            if ActionTypeHint.is_subclass_typehint(arg) and arg not in {object, type}:
+            if ActionTypeHint.is_subclass_typehint(arg, also_lists=True) and arg not in {object, type}:
                 add_subclasses(arg)
     else:
         add_subclasses(cls)

jsonargparse_tests/test_subclasses.py

@@ -1267,13 +1267,13 @@ def test_subclass_unresolved_parameters_name_clash(parser):
 def test_add_subclass_failure_not_a_class(parser):
     with pytest.raises(ValueError) as ctx:
         parser.add_subclass_arguments(NAMESPACE_OID, "oid")
-    ctx.match("Expected 'baseclass' argument to be a class or a tuple of classes")
+    ctx.match("Expected 'baseclass' to be a subclass type or a tuple of subclass types")
 
 
 def test_add_subclass_failure_empty_tuple(parser):
     with pytest.raises(ValueError) as ctx:
         parser.add_subclass_arguments((), "cls")
-    ctx.match("Expected 'baseclass' argument to be a class or a tuple of classes")
+    ctx.match("Expected 'baseclass' to be a subclass type or a tuple of subclass types")
 
 
 def test_add_subclass_lazy_default(parser):
@@ -1337,6 +1337,40 @@ def test_add_subclass_not_required_group(parser):
     assert init == Namespace()
 
 
+class ListUnionA:
+    def __init__(self, pa1: int):
+        self.pa1 = pa1
+
+
+class ListUnionB:
+    def __init__(self, pb1: str, pb2: float):
+        self.pb1 = pb1
+        self.pb2 = pb2
+
+
+def test_add_subclass_list_of_union(parser):
+    parser.add_argument("--config", action="config")
+    parser.add_subclass_arguments(
+        baseclass=(ListUnionA, ListUnionB, List[Union[ListUnionA, ListUnionB]]),
+        nested_key="subclass",
+    )
+    config = {
+        "subclass": [
+            {
+                "class_path": f"{__name__}.ListUnionB",
+                "init_args": {
+                    "pb1": "x",
+                    "pb2": 0.5,
+                },
+            }
+        ]
+    }
+    cfg = parser.parse_args([f"--config={config}"])
+    assert cfg.as_dict()["subclass"] == config["subclass"]
+    help_str = get_parser_help(parser)
+    assert "Show the help for the given subclass of {ListUnionA,ListUnionB}" in help_str
+
+
 # instance defaults tests