Repository: omni-us/jsonargparse
Pull request #186 (2022-10-28 05:27:31)

Support for Callable type with class as return type
## What does this PR do?

See the title and the changes in the README.rst file.

## Before submitting

- [x] Did you read the [contributing guideline](https://github.com/omni-us/jsonargparse/blob/master/CONTRIBUTING.rst)?
- [x] Did you update **the documentation**? (readme and public docstrings)
- [x] Did you write **unit tests** such that there is 100% coverage on related code? (required for bug fixes and new features)
- [x] Did you verify that new and existing **tests pass locally**?
- [x] Did you make sure that all changes preserve **backward compatibility**?
- [x] Did you update **the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)


sonarcloud[bot]: Kudos, SonarCloud Quality Gate passed!&nbsp; &nbsp; [![Quality Gate passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/passed-16px.png 'Quality Gate passed')](https://sonarcloud.io/dashboard?id=omni-us_jsonargparse&pullRequest=186)

[![Bug](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/bug-16px.png 'Bug')](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=BUG) [![A](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A-16px.png 'A')](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=BUG) [0 Bugs](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=BUG)  
[![Vulnerability](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/vulnerability-16px.png 'Vulnerability')](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=VULNERABILITY) [![A](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A-16px.png 'A')](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=VULNERABILITY) [0 Vulnerabilities](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=VULNERABILITY)  
[![Security Hotspot](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/security_hotspot-16px.png 'Security Hotspot')](https://sonarcloud.io/project/security_hotspots?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=SECURITY_HOTSPOT) [![A](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A-16px.png 'A')](https://sonarcloud.io/project/security_hotspots?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=SECURITY_HOTSPOT) [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=SECURITY_HOTSPOT)  
[![Code Smell](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/code_smell-16px.png 'Code Smell')](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=CODE_SMELL) [![A](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/RatingBadge/A-16px.png 'A')](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=CODE_SMELL) [2 Code Smells](https://sonarcloud.io/project/issues?id=omni-us_jsonargparse&pullRequest=186&resolved=false&types=CODE_SMELL)

[![No Coverage information](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/CoverageChart/NoCoverageInfo-16px.png 'No Coverage information')](https://sonarcloud.io/component_measures?id=omni-us_jsonargparse&pullRequest=186) No Coverage information  
[![0.0%](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/Duplications/3-16px.png '0.0%')](https://sonarcloud.io/component_measures?id=omni-us_jsonargparse&pullRequest=186&metric=new_duplicated_lines_density&view=list) [0.0% Duplication](https://sonarcloud.io/component_measures?id=omni-us_jsonargparse&pullRequest=186&metric=new_duplicated_lines_density&view=list)



Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 73faf37..e359451 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -20,6 +20,9 @@ Added
 - Type ``Any`` now parses and instantiates classes when given dict that follows
   subclass specification `pytorch-lightning#15115
   <https://github.com/Lightning-AI/lightning/issues/15115>`__.
+- Signature methods now accept skipping a number of positionals.
+- Callable type hint with return type a class can now be given a subclass which
+  produces a callable that returns an instance of the class.
 - Support for Python 3.11.
 
 Fixed
diff --git a/README.rst b/README.rst
index 25983ea..e310703 100644
--- a/README.rst
+++ b/README.rst
@@ -431,11 +431,13 @@ Some notes about this support are:
   classes in a config object. For more details see :ref:`sub-classes`.
 
 - ``Callable`` is supported by either giving a dot import path to a callable
-  object, or by giving a dict with a ``class_path`` and optionally ``init_args``
-  entries specifying a class that once instantiated is callable. Running
-  :py:meth:`.ArgumentParser.instantiate_classes` will instantiate the callable
-  classes. Currently the callable's arguments and return types are ignored and
-  not validated.
+  object or by giving a dict with a ``class_path`` and optionally ``init_args``
+  entries. The specified class must either instantiate into a callable or be a
+  subclass of the return type of the callable. For these cases running
+  :py:meth:`.ArgumentParser.instantiate_classes` will instantiate the class or
+  provide a function that returns the instance of the class. For more details
+  see :ref:`callable-type`. Currently the callable's argument and return types
+  are not validated.
 
 
 .. _restricted-numbers:
@@ -800,6 +802,91 @@ be achieved as follows:
     Namespace(dict={'key1': 'val1', 'key2': 'val2'})
 
 
+.. _callable-type:
+
+Callable type
+-------------
+
+When using ``Callable`` as type, the parser accepts several options. The first
+option is the import path of a callable object, for example:
+
+.. testsetup:: callable
+
+    import sys
+    from jsonargparse import ArgumentParser, ActionConfigFile
+    from typing import Callable, Iterable
+    parser = ArgumentParser()
+
+.. testcode:: callable
+
+    parser.add_argument('--callable', type=Callable)
+    parser.parse_args(['--callable=time.sleep'])
+
+A second option is a class that once instantiated becomes callable:
+
+.. code-block::
+
+    >>> class OffsetSum:
+    ...     def __init__(self, offset: int):
+    ...         self.offset = offset
+    ...     def __call__(self, value: int):
+    ...         return self.offset + value
+
+    >>> value = {
+    ...     'class_path': '__main__.OffsetSum',
+    ...     'init_args': {
+    ...         'offset': 3,
+    ...     }
+    ... }
+
+    >>> cfg = parser.parse_args(['--callable', str(value)])
+    >>> cfg.callable
+    Namespace(class_path='__main__.OffsetSum', init_args=Namespace(offset=3))
+    >>> init = parser.instantiate_classes(cfg)
+    >>> init.callable(5)
+    8
+
+The third option is only applicable when the type is a callable that has a class
+as return type. This is useful to support dependency injection for classes that
+require a parameter that is only available after injection. The parser does this
+automatically by providing a function that receives this parameter and returns
+the instance of the class.
+
+.. code-block::
+
+    >>> class Optimizer:
+    ...     pass
+
+    >>> class SGD(Optimizer):
+    ...     def __init__(self, params: Iterable, lr: float):
+    ...         self.params = params
+    ...         self.lr = lr
+
+    >>> value = {
+    ...     'class_path': 'SGD',
+    ...     'init_args': {
+    ...         'lr': 0.01,
+    ...     }
+    ... }
+
+    >>> parser.add_argument('--optimizer', type=Callable[[Iterable], Optimizer])
+    >>> cfg = parser.parse_args(['--optimizer', str(value)])
+    >>> cfg.optimizer
+    Namespace(class_path='__main__.SGD', init_args=Namespace(lr=0.01))
+    >>> init = parser.instantiate_classes(cfg)
+    >>> optim = init.optimizer([1, 2, 3])
+    >>> isinstance(optim, SGD)
+    True
+    >>> optim.params, optim.lr
+    ([1, 2, 3], 0.01)
+
+.. note::
+
+    When the ``Callable`` has a class return type, it is possible to specify the
+    ``class_path`` giving only its name if imported before parsing as explained
+    in :ref:`sub-classes-command-line`.
+
+
 .. _registering-types:
 
 Registering types
@@ -1500,6 +1587,7 @@ config file by specifying a path to it; and 3) by default sets a useful
     ``class_path`` and ``init_args`` if the corresponding parameter has type
     ``Any``, or when ``fail_untyped=False`` which defaults to type ``Any``.
 
+.. _sub-classes-command-line:
 
 Command line
 ------------
diff --git a/jsonargparse/signatures.py b/jsonargparse/signatures.py
index 6c357f4..c43202e 100644
--- a/jsonargparse/signatures.py
+++ b/jsonargparse/signatures.py
@@ -33,7 +33,7 @@ class SignatureArguments(LoggerProperty):
         as_group: bool = True,
         as_positional: bool = False,
         default: Optional[LazyInitBaseClass] = None,
-        skip: Optional[Set[str]] = None,
+        skip: Optional[Set[Union[str, int]]] = None,
         instantiate: bool = True,
         fail_untyped: bool = True,
         sub_configs: bool = False,
@@ -49,7 +49,7 @@ class SignatureArguments(LoggerProperty):
             as_group: Whether arguments should be added to a new argument group.
             as_positional: Whether to add required parameters as positional arguments.
             default: Default value used to override parameter defaults. Must be lazy_instance.
-            skip: Names of parameters that should be skipped.
+            skip: Names of parameters or number of positionals that should be skipped.
             instantiate: Whether the class group should be instantiated by :code:`instantiate_classes`.
             fail_untyped: Whether to raise exception if a required parameter does not have a type.
             sub_configs: Whether subclass type hints should be loadable from inner config file.
@@ -97,7 +97,7 @@ class SignatureArguments(LoggerProperty):
         nested_key: Optional[str] = None,
         as_group: bool = True,
         as_positional: bool = False,
-        skip: Optional[Set[str]] = None,
+        skip: Optional[Set[Union[str, int]]] = None,
         fail_untyped: bool = True,
         sub_configs: bool = False,
     ) -> List[str]:
@@ -111,7 +111,7 @@ class SignatureArguments(LoggerProperty):
             nested_key: Key for nested namespace.
             as_group: Whether arguments should be added to a new argument group.
             as_positional: Whether to add required parameters as positional arguments.
-            skip: Names of parameters that should be skipped.
+            skip: Names of parameters or number of positionals that should be skipped.
             fail_untyped: Whether to raise exception if a required parameter does not have a type.
             sub_configs: Whether subclass type hints should be loadable from inner config file.
 
@@ -145,7 +145,7 @@ class SignatureArguments(LoggerProperty):
         nested_key: Optional[str] = None,
         as_group: bool = True,
         as_positional: bool = False,
-        skip: Optional[Set[str]] = None,
+        skip: Optional[Set[Union[str, int]]] = None,
         fail_untyped: bool = True,
         sub_configs: bool = False,
     ) -> List[str]:
@@ -158,7 +158,7 @@ class SignatureArguments(LoggerProperty):
             nested_key: Key for nested namespace.
             as_group: Whether arguments should be added to a new argument group.
             as_positional: Whether to add required parameters as positional arguments.
-            skip: Names of parameters that should be skipped.
+            skip: Names of parameters or number of positionals that should be skipped.
             fail_untyped: Whether to raise exception if a required parameter does not have a type.
             sub_configs: Whether subclass type hints should be loadable from inner config file.
 
@@ -191,7 +191,7 @@ class SignatureArguments(LoggerProperty):
         nested_key: Optional[str],
         as_group: bool = True,
         as_positional: bool = False,
-        skip: Optional[Set[str]] = None,
+        skip: Optional[Set[Union[str, int]]] = None,
         fail_untyped: bool = True,
         sub_configs: bool = False,
         instantiate: bool = True,
@@ -205,7 +205,7 @@ class SignatureArguments(LoggerProperty):
             nested_key: Key for nested namespace.
             as_group: Whether arguments should be added to a new argument group.
             as_positional: Whether to add required parameters as positional arguments.
-            skip: Names of parameters that should be skipped.
+            skip: Names of parameters or number of positionals that should be skipped.
             fail_untyped: Whether to raise exception if a required parameter does not have a type.
             sub_configs: Whether subclass type hints should be loadable from inner config file.
             instantiate: Whether the class group should be instantiated by :code:`instantiate_classes`.
@@ -218,6 +218,14 @@ class SignatureArguments(LoggerProperty):
         """
         params = get_signature_parameters(function_or_class, method_name, logger=self.logger)
 
+        skip_positionals = [s for s in (skip or []) if isinstance(s, int)]
+        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}')
+            names = {p.name for p in params[:skip_positionals[0]]}
+            params = params[skip_positionals[0]:]
+            self.logger.debug(f'Skipping parameters {names} because {skip_positionals[0]} positionals requested to be skipped.')
+
         ## 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
@@ -238,7 +246,7 @@ class SignatureArguments(LoggerProperty):
                 nested_key,
                 param,
                 added_args,
-                skip,
+                skip={s for s in (skip or []) if isinstance(s, str)},
                 fail_untyped=fail_untyped,
                 sub_configs=sub_configs,
                 linked_targets=linked_targets,
diff --git a/jsonargparse/typehints.py b/jsonargparse/typehints.py
index 5c5ca9c..da37a31 100644
--- a/jsonargparse/typehints.py
+++ b/jsonargparse/typehints.py
@@ -433,12 +433,15 @@ class ActionTypeHint(Action):
 
 
     @staticmethod
-    def get_class_parser(val_class, sub_add_kwargs=None):
+    def get_class_parser(val_class, sub_add_kwargs=None, skip_args=0):
         from .core import ArgumentParser
         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)
         parser = ArgumentParser(error_handler=None)
-        parser.add_class_arguments(val_class, **(sub_add_kwargs or {}))
+        parser.add_class_arguments(val_class, **kwargs)
         return parser
 
 
@@ -446,7 +449,7 @@ class ActionTypeHint(Action):
         extra = ''
         if (
             self.is_subclass_typehint(self, all_subtypes=False) or
-            get_typehint_origin(self._typehint) in {Type, type}
+            get_typehint_origin(self._typehint) in callable_origin_types.union({Type, type})
         ):
             class_paths = get_all_subclass_paths(self._typehint)
             if class_paths:
@@ -659,11 +662,12 @@ def adapt_typehints(val, typehint, serialize=False, instantiate_classes=False, p
                     val = subclass_spec_as_namespace(val, prev_val)
                     if not is_subclass_spec(val):
                         raise ImportError(f'Dict must include a class_path and optionally init_args, but got {val_input}')
+                    val, partial_classes, num_partial_args = get_partial_callable_classes(typehint, val)
                     val_class = import_object(val['class_path'])
-                    if not (inspect.isclass(val_class) and callable_instances(val_class)):
-                        raise ImportError(f'{val["class_path"]!r} is not a callable class.')
+                    if inspect.isclass(val_class) and not (partial_classes or callable_instances(val_class)):
+                        raise ImportError(f'Expected {val["class_path"]!r} to be a class that instantiates into callable or a subclass of {partial_classes}.')
                     val['class_path'] = get_import_path(val_class)
-                    val = adapt_class_type(val, False, instantiate_classes, sub_add_kwargs)
+                    val = adapt_class_type(val, False, instantiate_classes, sub_add_kwargs, skip_args=num_partial_args)
             except (ImportError, AttributeError, ParserError) as ex:
                 raise ValueError(f'Type {typehint} expects a function or a callable class: {ex}') from ex
 
@@ -742,6 +746,26 @@ def subclass_spec_as_namespace(val, prev_val=None):
     return val
 
 
+def get_partial_callable_classes(callable_type, subclass_spec):
+    partial_classes = False
+    num_partial_args = 0
+    subtypes = getattr(callable_type, '__args__', None)
+    if subtypes and len(subtypes) > 1:
+        return_type = subtypes[-1]
+        return_origin = get_typehint_origin(return_type)
+        subclass_types = tuple(
+            t for t in (return_type.__args__ if return_origin == Union else [return_type])
+            if ActionTypeHint.is_subclass_typehint(t)
+        )
+        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 = deepcopy(subclass_spec)
+            subclass_spec['class_path'] = get_import_path(class_type)
+            partial_classes = True
+            num_partial_args = len(subtypes)-1
+    return subclass_spec, partial_classes, num_partial_args
+
+
 def get_all_subclass_paths(cls: Type) -> List[str]:
     subclass_list = []
 
@@ -760,6 +784,11 @@ def get_all_subclass_paths(cls: Type) -> List[str]:
         for subclass in cl.__subclasses__() if hasattr(cl, '__subclasses__') else []:
             add_subclasses(subclass)
 
+    if get_typehint_origin(cls) in callable_origin_types:
+        if len(getattr(cls, '__args__', [])) < 2:
+            return subclass_list
+        cls = cls.__args__[-1]
+
     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}:
@@ -827,11 +856,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):
+def adapt_class_type(value, serialize, instantiate_classes, sub_add_kwargs, prev_val=None, skip_args=0):
     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)
+    parser = ActionTypeHint.get_class_parser(val_class, sub_add_kwargs, skip_args=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', []):
@@ -860,6 +889,10 @@ def adapt_class_type(value, serialize, instantiate_classes, sub_add_kwargs, prev
             if init_args:
                 value['init_args'] = init_args
             return value
+        if skip_args:
+            def partial_instance(*args):
+                return val_class(*args, **{**init_args, **dict_kwargs})
+            return partial_instance
         return val_class(**{**init_args, **dict_kwargs})
 
     prev_init_args = prev_val.get('init_args') if isinstance(prev_val, Namespace) else None


Test patch:
diff --git a/jsonargparse_tests/test_signatures.py b/jsonargparse_tests/test_signatures.py
index 9ee7ba1..537c28c 100755
--- a/jsonargparse_tests/test_signatures.py
+++ b/jsonargparse_tests/test_signatures.py
@@ -630,13 +630,20 @@ class SignaturesTests(unittest.TestCase):
                  a4: int = 4):
             return a1
 
-        parser = ArgumentParser()
-        parser.add_function_arguments(func, skip={'a2', 'a4'})
+        with self.subTest('parameter names'):
+            parser = ArgumentParser()
+            added_args = parser.add_function_arguments(func, skip={'a2', 'a4'})
+            self.assertEqual(added_args, ['a1', 'a3'])
 
-        for key in ['a1', 'a3']:
-            self.assertIsNotNone(_find_action(parser, key), key+' should be in parser but is not')
-        for key in ['a2', 'a4']:
-            self.assertIsNone(_find_action(parser, key), key+' should not be in parser but is')
+        with self.subTest('positionals and names'):
+            parser = ArgumentParser()
+            added_args = parser.add_function_arguments(func, skip={1, 'a3'})
+            self.assertEqual(added_args, ['a2', 'a4'])
+
+        with self.subTest('invalid skip'):
+            parser = ArgumentParser()
+            with self.assertRaises(ValueError):
+                parser.add_function_arguments(func, skip={1, 2})
 
 
     def test_skip_within_subclass_type(self):
diff --git a/jsonargparse_tests/test_typehints.py b/jsonargparse_tests/test_typehints.py
index 4b374c1..502d59d 100755
--- a/jsonargparse_tests/test_typehints.py
+++ b/jsonargparse_tests/test_typehints.py
@@ -6,6 +6,7 @@ import pathlib
 import random
 import re
 import sys
+import time
 import unittest
 import uuid
 import warnings
@@ -601,9 +602,14 @@ class TypeHintsTests(unittest.TestCase):
 
     def test_Callable_with_function_path(self):
         parser = ArgumentParser(error_handler=None)
-        parser.add_argument('--callable', type=Callable, default=lazy_instance)
+        parser.add_argument('--callable', type=Callable, default=time.time)
         parser.add_argument('--list', type=List[Callable])
 
+        cfg = parser.get_defaults()
+        self.assertEqual(time.time, cfg.callable)
+        self.assertEqual(parser.dump(cfg), 'callable: time.time\n')
+        cfg = parser.parse_args(['--callable=random.randint'])
+        self.assertEqual(random.randint, cfg.callable)
         cfg = parser.parse_args(['--callable=jsonargparse.CLI'])
         self.assertEqual(CLI, cfg.callable)
         self.assertEqual(parser.dump(cfg), 'callable: jsonargparse.CLI\n')
@@ -612,7 +618,7 @@ class TypeHintsTests(unittest.TestCase):
 
         out = StringIO()
         parser.print_help(out)
-        self.assertIn('(type: Callable, default: jsonargparse.lazy_instance)', out.getvalue())
+        self.assertIn('(type: Callable, default: time.time)', out.getvalue())
 
 
     def test_Callable_with_class_path(self):
@@ -700,6 +706,101 @@ class TypeHintsTests(unittest.TestCase):
             self.assertEqual(my_func_2, cfg.callable)  # Currently callable types are ignored
 
 
+    def test_typed_callable_with_return_type_class(self):
+        class Optimizer:
+            def __init__(self, params: List[float], lr: float = 1e-3):
+                self.params = params
+                self.lr = lr
+
+        class SGD(Optimizer):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+
+        class Adam(Optimizer):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+
+        value = {
+            'class_path': 'Adam',
+            'init_args': {
+                'lr': 0.01,
+            }
+        }
+
+        with mock_module(Optimizer, SGD, Adam) as module:
+            parser = ArgumentParser(error_handler=None)
+            parser.add_argument('--optimizer', type=Callable[[List[float]], Optimizer], default=SGD)
+
+            cfg = parser.get_defaults()
+            init = parser.instantiate_classes(cfg)
+            optim = init.optimizer([0.1, 2,3])
+            self.assertIsInstance(optim, SGD)
+            self.assertEqual(optim.params, [0.1, 2,3])
+            self.assertEqual(optim.lr, 1e-3)
+
+            cfg = parser.parse_args(['--optimizer', str(value)])
+            self.assertEqual(cfg.optimizer.class_path, f'{module}.Adam')
+            self.assertEqual(cfg.optimizer.init_args, Namespace(lr=0.01))
+            init = parser.instantiate_classes(cfg)
+            optim = init.optimizer([4.5, 6.7])
+            self.assertIsInstance(optim, Adam)
+            self.assertEqual(optim.params, [4.5, 6.7])
+            self.assertEqual(optim.lr, 0.01)
+
+            help_str = StringIO()
+            parser.print_help(help_str)
+            for name in ['Optimizer', 'SGD', 'Adam']:
+                self.assertIn(f'{module}.{name}', help_str.getvalue())
+
+
+    def test_typed_callable_with_return_type_union_of_classes(self):
+        class Optimizer:
+            pass
+
+        class StepLR:
+            def __init__(self, optimizer: Optimizer, last_epoch: int = -1):
+                self.optimizer = optimizer
+                self.last_epoch = last_epoch
+
+        class ReduceLROnPlateau:
+            def __init__(self, optimizer: Optimizer, monitor: str):
+                self.optimizer = optimizer
+                self.monitor = monitor
+
+        optim = Optimizer()
+        value = {
+            'class_path': 'ReduceLROnPlateau',
+            'init_args': {
+                'monitor': 'loss',
+            }
+        }
+
+        with mock_module(StepLR, ReduceLROnPlateau) as module:
+            parser = ArgumentParser(error_handler=None)
+            parser.add_argument('--scheduler', type=Callable[[Optimizer], Union[StepLR, ReduceLROnPlateau]], default=StepLR)
+
+            cfg = parser.get_defaults()
+            init = parser.instantiate_classes(cfg)
+            sched = init.scheduler(optim)
+            self.assertIsInstance(sched, StepLR)
+            self.assertIs(sched.optimizer, optim)
+            self.assertEqual(sched.last_epoch, -1)
+
+            cfg = parser.parse_args(['--scheduler', str(value)])
+            self.assertEqual(cfg.scheduler.class_path, f'{module}.ReduceLROnPlateau')
+            self.assertEqual(cfg.scheduler.init_args, Namespace(monitor='loss'))
+            init = parser.instantiate_classes(cfg)
+            sched = init.scheduler(optim)
+            self.assertIsInstance(sched, ReduceLROnPlateau)
+            self.assertIs(sched.optimizer, optim)
+            self.assertEqual(sched.monitor, 'loss')
+
+            help_str = StringIO()
+            parser.print_help(help_str)
+            for name in ['StepLR', 'ReduceLROnPlateau']:
+                self.assertIn(f'{module}.{name}', help_str.getvalue())
+
+
     def test_class_type(self):
         parser = ArgumentParser(error_handler=None)
         parser.add_argument('--op', type=Optional[List[Calendar]])