Add support callable protocols for instance factory dependency injection (#755).

Author: mauvilsa

CHANGELOG.rst

@@ -19,6 +19,8 @@ Added
 ^^^^^
 - Support for Python 3.14 (`#753
   <https://github.com/omni-us/jsonargparse/pull/753>`__).
+- Support callable protocols for instance factory dependency injection (`#758
+  <https://github.com/omni-us/jsonargparse/pull/758>`__).
 
 Changed
 ^^^^^^^

jsonargparse/_actions.py

@@ -350,7 +350,6 @@ def __init__(self, typehint=None, **kwargs):
         if typehint is not None:
             self._typehint = typehint
         else:
-            self._typehint = kwargs.pop("_typehint")
             self.update_init_kwargs(kwargs)
             super().__init__(**kwargs)
 
@@ -363,13 +362,14 @@ def update_init_kwargs(self, kwargs):
             is_protocol,
         )
 
-        typehint = get_unaliased_type(get_optional_arg(self._typehint))
+        typehint = get_unaliased_type(get_optional_arg(kwargs.pop("_typehint")))
         if get_typehint_origin(typehint) is not Union:
             assert "nargs" not in kwargs
             kwargs["nargs"] = "?"
-        self._basename = iter_to_set_str(get_subclass_names(self._typehint, callable_return=True))
+        self._typehint = typehint
+        self._basename = iter_to_set_str(get_subclass_names(typehint, callable_return=True))
         self._baseclasses = get_subclass_types(typehint, callable_return=True)
-        assert self._baseclasses
+        assert self._baseclasses and all(isinstance(b, type) for b in self._baseclasses)
 
         self._kind = "subclass of"
         if any(is_protocol(b) for b in self._baseclasses):
@@ -391,28 +391,28 @@ def __call__(self, *args, **kwargs):
 
     def print_help(self, call_args):
         from ._typehints import (
-            ActionTypeHint,
-            get_optional_arg,
-            get_unaliased_type,
+            adapt_partial_callable_class,
             implements_protocol,
             resolve_class_path_by_name,
         )
 
         parser, _, value, option_string = call_args
         try:
-            typehint = get_unaliased_type(get_optional_arg(self._typehint))
             if self.nargs == "?" and value is None:
-                val_class = typehint
+                val_class = self._typehint
             else:
-                val_class = import_object(resolve_class_path_by_name(typehint, value))
+                val_class = import_object(resolve_class_path_by_name(self._baseclasses, value))
         except Exception as ex:
             raise TypeError(f"{option_string}: {ex}") from ex
+
         if not any(is_subclass(val_class, b) or implements_protocol(val_class, b) for b in self._baseclasses):
             raise TypeError(f'{option_string}: Class "{value}" is not a {self._kind} {self._basename}')
         dest = re.sub("\\.help$", "", self.dest)
         subparser = type(parser)(description=f"Help for {option_string}={get_import_path(val_class)}")
-        if ActionTypeHint.is_callable_typehint(typehint) and hasattr(typehint, "__args__"):
-            self.sub_add_kwargs["skip"] = {max(0, len(typehint.__args__) - 1)}
+        val = Namespace(class_path=get_import_path(val_class))
+        _, partial_skip_args = adapt_partial_callable_class(self._typehint, val)
+        if partial_skip_args:
+            self.sub_add_kwargs["skip"] = partial_skip_args
         subparser.add_class_arguments(val_class, dest, **self.sub_add_kwargs)
         subparser._inner_parser = True
         remove_actions(subparser, (_HelpAction, _ActionPrintConfig, _ActionConfigLoad))

jsonargparse/_signatures.py

@@ -267,7 +267,7 @@ def _add_signature_arguments(
         """
         params = get_signature_parameters(function_or_class, method_name, logger=self.logger)
 
-        skip_positionals = [s for s in (skip or []) if isinstance(s, int)]
+        skip_positionals = [s for s in (skip or []) if isinstance(s, int) and s != 0]
         if skip_positionals:
             if len(skip_positionals) > 1 or any(p <= 0 for p in skip_positionals):
                 raise ValueError(f"Unexpected number of positionals to skip: {skip_positionals}")

jsonargparse/_typehints.py

@@ -271,14 +271,14 @@ def normalize_default(self, default):
             default = default.name
         elif is_callable_type(self._typehint) and callable(default) and not inspect.isclass(default):
             default = get_import_path(default)
+        elif ActionTypeHint.is_return_subclass_typehint(self._typehint) and inspect.isclass(default):
+            default = {"class_path": get_import_path(default)}
         elif is_subclass_type and not allow_default_instance.get():
             from ._parameter_resolvers import UnknownDefault
 
             default_type = type(default)
             if not is_subclass(default_type, UnknownDefault) and self.is_subclass_typehint(default_type):
                 raise ValueError("Subclass types require as default either a dict with class_path or a lazy instance.")
-        elif ActionTypeHint.is_return_subclass_typehint(self._typehint) and inspect.isclass(default):
-            default = {"class_path": get_import_path(default)}
         return default
 
     @staticmethod
@@ -352,7 +352,7 @@ def is_subclass_typehint(typehint, all_subtypes=True, also_lists=False):
     def is_return_subclass_typehint(typehint):
         typehint = get_unaliased_type(get_optional_arg(get_unaliased_type(typehint)))
         typehint_origin = get_typehint_origin(typehint)
-        if typehint_origin in callable_origin_types:
+        if typehint_origin in callable_origin_types or is_instance_factory_protocol(typehint):
             return_type = get_callable_return_type(typehint)
             if ActionTypeHint.is_subclass_typehint(return_type):
                 return True
@@ -626,15 +626,18 @@ def instantiate_classes(self, value):
         return value if islist else value[0]
 
     @staticmethod
-    def get_class_parser(val_class, sub_add_kwargs=None, skip_args=0):
+    def get_class_parser(val_class, sub_add_kwargs=None, skip_args=None):
         if isinstance(val_class, str):
             val_class = import_object(val_class)
         kwargs = dict(sub_add_kwargs) if sub_add_kwargs else {}
         if skip_args:
-            kwargs.setdefault("skip", set()).add(skip_args)
+            kwargs.setdefault("skip", set()).update(skip_args)
         if is_subclass_spec(kwargs.get("default")):
             kwargs["default"] = kwargs["default"].get("init_args")
         parser = parent_parser.get()
+        from ._core import ArgumentParser
+
+        assert isinstance(parser, ArgumentParser)
         parser = type(parser)(exit_on_error=False, logger=parser.logger, parser_mode=parser.parser_mode)
         remove_actions(parser, (ActionConfigFile, _ActionPrintConfig))
         if inspect.isclass(val_class) or inspect.isclass(get_typehint_origin(val_class)):
@@ -658,11 +661,10 @@ def get_class_parser(val_class, sub_add_kwargs=None, skip_args=0):
     def extra_help(self):
         extra = ""
         typehint = get_optional_arg(self._typehint)
-        if self.is_subclass_typehint(typehint, all_subtypes=False) or get_typehint_origin(
-            typehint
-        ) in callable_origin_types.union({Type, type}):
-            if self.is_callable_typehint(typehint) and getattr(typehint, "__args__", None):
-                typehint = get_callable_return_type(get_optional_arg(typehint))
+        typehint = get_callable_return_type(typehint) or typehint
+        if get_typehint_origin(typehint) is type:
+            typehint = typehint.__args__[0]
+        if self.is_subclass_typehint(typehint, all_subtypes=False):
             class_paths = get_all_subclass_paths(typehint)
             if class_paths:
                 extra = ", known subclasses: " + ", ".join(class_paths)
@@ -967,11 +969,15 @@ def adapt_typehints(
         val = adapt_typehints(val, subtypehints[0], **adapt_kwargs)
 
     # Callable
-    elif typehint_origin in callable_origin_types or typehint in callable_origin_types:
+    elif (
+        typehint_origin in callable_origin_types
+        or typehint in callable_origin_types
+        or is_instance_factory_protocol(typehint, logger)
+    ):
         if serialize:
             if is_subclass_spec(val):
-                val, _, num_partial_args = adapt_partial_callable_class(typehint, val)
-                val = adapt_class_type(val, True, False, sub_add_kwargs, skip_args=num_partial_args)
+                val, partial_skip_args = adapt_partial_callable_class(typehint, val)
+                val = adapt_class_type(val, True, False, sub_add_kwargs, partial_skip_args=partial_skip_args)
             else:
                 val = object_path_serializer(val)
         else:
@@ -1000,21 +1006,21 @@ def adapt_typehints(
                         raise ImportError(
                             f"Dict must include a class_path and optionally init_args, but got {val_input}"
                         )
-                    val, partial_classes, num_partial_args = adapt_partial_callable_class(typehint, val)
+                    val, partial_skip_args = adapt_partial_callable_class(typehint, val)
                     val_class = import_object(val["class_path"])
-                    if inspect.isclass(val_class) and not (partial_classes or callable_instances(val_class)):
+                    if inspect.isclass(val_class) and not (partial_skip_args or callable_instances(val_class)):
+                        base_type = get_callable_return_type(typehint) or typehint
                         raise ImportError(
                             f"Expected '{val['class_path']}' to be a class that instantiates into callable "
-                            f"or a subclass of {partial_classes}."
+                            f"or a subclass of {base_type}."
                         )
                     val["class_path"] = get_import_path(val_class)
                     val = adapt_class_type(
                         val,
                         False,
                         instantiate_classes,
                         sub_add_kwargs,
-                        skip_args=num_partial_args,
-                        partial_classes=partial_classes,
+                        partial_skip_args=partial_skip_args,
                         prev_val=prev_val,
                     )
             except (ImportError, AttributeError, ArgumentError) as ex:
@@ -1172,6 +1178,15 @@ def is_instance_or_supports_protocol(value, class_type):
     return isinstance(value, class_type)
 
 
+def is_instance_factory_protocol(class_type, logger=None):
+    if not is_protocol(class_type) or not callable_instances(class_type):
+        return False
+    from ._postponed_annotations import get_return_type
+
+    return_type = get_return_type(class_type.__call__, logger)
+    return ActionTypeHint.is_subclass_typehint(return_type)
+
+
 def is_subclass_spec(val):
     is_class = isinstance(val, (dict, Namespace)) and "class_path" in val
     if is_class:
@@ -1214,9 +1229,14 @@ def subclass_spec_as_namespace(val, prev_val=None):
 
 def get_callable_return_type(typehint):
     return_type = None
-    args = getattr(typehint, "__args__", None)
-    if isinstance(args, tuple) and len(args) > 0:
-        return_type = args[-1]
+    if is_instance_factory_protocol(typehint):
+        from ._postponed_annotations import get_return_type
+
+        return_type = get_return_type(typehint.__call__)
+    elif get_typehint_origin(typehint) in callable_origin_types:
+        args = getattr(typehint, "__args__", None)
+        if isinstance(args, tuple) and len(args) > 0:
+            return_type = args[-1]
     return return_type
 
 
@@ -1238,7 +1258,7 @@ def yield_subclass_types(typehint, also_lists=False, callable_return=False):
         return
     typehint = get_unaliased_type(get_optional_arg(get_unaliased_type(typehint)))
     typehint_origin = get_typehint_origin(typehint)
-    if callable_return and typehint_origin in callable_origin_types:
+    if callable_return and (typehint_origin in callable_origin_types or is_instance_factory_protocol(typehint)):
         return_type = get_callable_return_type(typehint)
         if return_type:
             k = {"also_lists": also_lists, "callable_return": callable_return}
@@ -1261,18 +1281,26 @@ def get_subclass_names(typehint, callable_return=False):
 
 
 def adapt_partial_callable_class(callable_type, subclass_spec):
-    partial_classes = False
-    num_partial_args = 0
+    partial_skip_args = None
     return_type = get_callable_return_type(callable_type)
     if return_type:
         subclass_types = get_subclass_types(return_type)
         class_type = import_object(resolve_class_path_by_name(return_type, subclass_spec.class_path))
         if subclass_types and is_subclass(class_type, subclass_types):
             subclass_spec = subclass_spec.clone()
             subclass_spec["class_path"] = get_import_path(class_type)
-            partial_classes = True
-            num_partial_args = len(callable_type.__args__) - 1
-    return subclass_spec, partial_classes, num_partial_args
+            if is_protocol(callable_type):
+                from ._parameter_resolvers import get_signature_parameters
+
+                params = get_signature_parameters(callable_type, "__call__")
+                partial_skip_args = set()
+                positionals = [p for p in params if "POSITIONAL_ONLY" in str(p.kind)]
+                if positionals:
+                    partial_skip_args.add(len(positionals))
+                partial_skip_args.update(p.name for p in params if "POSITIONAL_ONLY" not in str(p.kind))
+            else:
+                partial_skip_args = {len(callable_type.__args__) - 1}
+    return subclass_spec, partial_skip_args
 
 
 def get_all_subclass_paths(cls: Type) -> List[str]:
@@ -1318,9 +1346,15 @@ def add_subclasses(cl):
     return subclass_list
 
 
-def resolve_class_path_by_name(cls: Type, name: str) -> str:
+def resolve_class_path_by_name(cls: Union[Type, Tuple[Type]], name: str) -> str:
     class_path = name
     if "." not in class_path:
+        if isinstance(cls, tuple):
+            for cls_n in cls:
+                class_path = resolve_class_path_by_name(cls_n, name)
+                if "." in class_path:
+                    break
+            return class_path
         subclass_dict = defaultdict(list)
         for subclass in get_all_subclass_paths(cls):
             subclass_name = subclass.rsplit(".", 1)[1]
@@ -1376,13 +1410,11 @@ def discard_init_args_on_class_path_change(parser_or_action, prev_val, value):
             )
 
 
-def adapt_class_type(
-    value, serialize, instantiate_classes, sub_add_kwargs, prev_val=None, skip_args=0, partial_classes=False
-):
+def adapt_class_type(value, serialize, instantiate_classes, sub_add_kwargs, prev_val=None, partial_skip_args=None):
     prev_val = subclass_spec_as_namespace(prev_val)
     value = subclass_spec_as_namespace(value)
     val_class = import_object(value.class_path)
-    parser = ActionTypeHint.get_class_parser(val_class, sub_add_kwargs, skip_args=skip_args)
+    parser = ActionTypeHint.get_class_parser(val_class, sub_add_kwargs, skip_args=partial_skip_args)
 
     # No need to re-create the linked arg but just "inform" the corresponding parser actions that it exists upstream.
     for target in sub_add_kwargs.get("linked_targets", []):
@@ -1415,7 +1447,7 @@ def adapt_class_type(
 
         instantiator_fn = get_class_instantiator()
 
-        if partial_classes:
+        if partial_skip_args:
             return partial(
                 instantiator_fn,
                 val_class,

jsonargparse_tests/test_typehints.py

@@ -21,6 +21,7 @@
     Literal,
     Mapping,
     Optional,
+    Protocol,
     Sequence,
     Set,
     Tuple,
@@ -1058,6 +1059,72 @@ def test_callable_args_return_type_class(parser, subtests):
         assert "--optimizer.params" not in help_str
 
 
+class OptimizerFactory(Protocol):
+    def __call__(self, params: List[float]) -> Optimizer: ...
+
+
+class DifferentParamsOrder(Optimizer):
+    def __init__(self, lr: float, params: List[float], momentum: float = 0.0):
+        super().__init__(lr=lr, params=params, momentum=momentum)
+
+
+def test_callable_protocol_instance_factory(parser, subtests):
+    parser.add_argument("--optimizer", type=OptimizerFactory, default=SGD)
+
+    with subtests.test("default"):
+        cfg = parser.get_defaults()
+        assert cfg.optimizer.class_path == f"{__name__}.SGD"
+        init = parser.instantiate_classes(cfg)
+        optimizer = init.optimizer(params=[1, 2])
+        assert isinstance(optimizer, SGD)
+        assert optimizer.params == [1, 2]
+        assert optimizer.lr == 1e-3
+        assert optimizer.momentum == 0.0
+
+    with subtests.test("parse dict"):
+        value = {
+            "class_path": "Adam",
+            "init_args": {
+                "lr": 0.01,
+                "momentum": 0.9,
+            },
+        }
+        cfg = parser.parse_args([f"--optimizer={json.dumps(value)}"])
+        init = parser.instantiate_classes(cfg)
+        optimizer = init.optimizer(params=[3, 2, 1])
+        assert isinstance(optimizer, Adam)
+        assert optimizer.params == [3, 2, 1]
+        assert optimizer.lr == 0.01
+        assert optimizer.momentum == 0.9
+
+    with subtests.test("params order"):
+        value = {
+            "class_path": "DifferentParamsOrder",
+            "init_args": {
+                "lr": 0.1,
+                "momentum": 0.8,
+            },
+        }
+        cfg = parser.parse_args([f"--optimizer={json.dumps(value)}"])
+        init = parser.instantiate_classes(cfg)
+        optimizer = init.optimizer(params=[3, 2])
+        assert isinstance(optimizer, DifferentParamsOrder)
+        assert optimizer.params == [3, 2]
+        assert optimizer.lr == 0.1
+        assert optimizer.momentum == 0.8
+        dump = parser.dump(cfg)
+        assert json_or_yaml_load(dump) == cfg.as_dict()
+
+    with subtests.test("help"):
+        help_str = get_parser_help(parser)
+        assert "--optimizer.help" in help_str
+        assert "Show the help for the given subclass or implementer of protocol {Optimizer,OptimizerFactory" in help_str
+        help_str = get_parse_args_stdout(parser, [f"--optimizer.help={__name__}.DifferentParamsOrder"])
+        assert f"Help for --optimizer.help={__name__}.DifferentParamsOrder" in help_str
+        assert "--optimizer.lr" in help_str
+        assert "--optimizer.params" not in help_str
+
+
 def test_optional_callable_return_type_help(parser):
     parser.add_argument("--optimizer", type=Optional[Callable[[List[float]], Optimizer]])
     help_str = get_parser_help(parser)