Support type[SomeTypedDict], ModuleType, UnionType and GenericAlias as type hints

CHANGELOG.rst

@@ -29,6 +29,20 @@ Added
 - ``shtab`` completion scripts now include file and directory completions for
   arguments typed as pydantic's ``FilePath`` and ``DirectoryPath`` (`#943
   <https://github.com/mauvilsa/jsonargparse/pull/943>`__).
+- Support ``type[SomeTypedDict]`` such that the given class is accepted when it
+  is structurally compatible, i.e. it has all the keys of the expected
+  ``TypedDict``, with the same types and requiredness (`#945
+  <https://github.com/mauvilsa/jsonargparse/pull/945>`__).
+- Support ``types.ModuleType`` as a type. The value is the import path of a
+  module, which on parse is validated to be importable, and on ``instantiate``
+  is replaced by the imported module object (`#945
+  <https://github.com/mauvilsa/jsonargparse/pull/945>`__).
+- Support ``types.UnionType`` and ``types.GenericAlias`` as types, often seen in
+  third party libraries in unions such as ``type | UnionType | dict``. The value
+  is a string with a type expression, e.g. ``"int | str"`` and ``"list[int]"``.
+  Previously adding an argument with these types failed with ``TypeError:
+  'member_descriptor' object is not iterable`` (`#945
+  <https://github.com/mauvilsa/jsonargparse/pull/945>`__).
 
 Fixed
 ^^^^^

DOCUMENTATION.rst

@@ -532,7 +532,13 @@ Some notes about this support are:
   fine-grained specification of required/optional ``TypedDict`` keys.
   ``Unpack`` is supported with ``TypedDict`` for more precise ``**kwargs``
   typing as described in PEP `692 <https://peps.python.org/pep-0692/>`__.
-  For more details see :ref:`dict-items`.
+  For more details see :ref:`dict-items`. A ``TypedDict`` can also be used as
+  the argument of ``type``, e.g. ``type[SomeTypedDict]``, in which case the
+  value is an import path to a class. Since ``TypedDict`` classes don't support
+  ``issubclass``, the given class is accepted when it is structurally
+  compatible, as specified in PEP `589 <https://peps.python.org/pep-0589/>`__,
+  i.e. it has all the keys of the expected ``TypedDict``, with the same types
+  and requiredness.
 
 - ``tuple``, ``set``, ``frozenset`` and ``MutableSet`` are supported even though
   they can't be represented in JSON distinguishable from a list. Each ``tuple``
@@ -591,6 +597,15 @@ Some notes about this support are:
   see :ref:`callable-type`. Currently the callable's argument and return types
   are not validated.
 
+- ``types.ModuleType`` is supported by giving the dot import path of a module,
+  and on ``instantiate`` is replaced by the imported module object.
+
+- ``types.UnionType`` and ``types.GenericAlias``, commonly found in third party
+  libraries in unions such as ``type | UnionType | dict``, are supported by
+  giving a string with a type expression, e.g. ``"int | str"`` and
+  ``"list[int]"``. The expression is resolved without evaluating code, so its
+  names must be builtins, ``typing`` names or dot import paths.
+
 - ``TypeAliasType`` is supported with values parsed as the aliased type and the
   alias shown as the argument type in help.
 

jsonargparse/_link_arguments.py

@@ -56,7 +56,7 @@ def find_subclass_action_or_class_group(
     from ._typehints import ActionTypeHint
 
     action = find_parent_action(parser, key, exclude=exclude)
-    if ActionTypeHint.is_subclass_typehint(action):
+    if ActionTypeHint.is_subclass_typehint(action) or ActionTypeHint.is_module_typehint(action):
         return action
     key_set = {key, split_key_leaf(key)[0]}
     for group in parser._action_groups:

jsonargparse/_typehints.py

@@ -1,9 +1,12 @@
 """Action to support type hints."""
 
+import ast
+import builtins
 import inspect
 import os
 import re
 import sys
+import typing
 from argparse import ArgumentError
 from collections import OrderedDict, abc, defaultdict, deque
 from contextlib import contextmanager, suppress
@@ -12,7 +15,8 @@
 from enum import Enum
 from functools import partial
 from importlib import import_module
-from types import FunctionType, MappingProxyType
+from importlib.util import find_spec
+from types import FunctionType, GenericAlias, MappingProxyType, ModuleType, UnionType
 from typing import (
     Any,
     Callable,
@@ -151,6 +155,9 @@
     OrderedDict,
     Callable,
     abc.Callable,
+    ModuleType,
+    UnionType,
+    GenericAlias,
     NotRequired,
     Required,
     Unpack,
@@ -344,7 +351,7 @@
             self._supports_append = self.supports_append(self._typehint)
             self.default = self.normalize_default(self.default)
 
     def normalize_default(self, default):
         from ._signatures import convert_to_dict, is_convertible_to_dict
 
         is_subclass_type = self.is_subclass_typehint(self._typehint, all_subtypes=False)
@@ -357,6 +364,8 @@
             default.class_path = normalize_import_path(default.class_path, self._typehint)
         elif is_enum_type(self._typehint) and isinstance(default, Enum):
             default = default.name
+        elif is_module_type(self._typehint) and isinstance(default, ModuleType):
+            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):
@@ -461,6 +470,11 @@
             return True
         return False
 
+    @staticmethod
+    def is_module_typehint(typehint):
+        typehint = typehint_from_action(typehint)
+        return typehint is not None and is_module_type(typehint)
+
     @staticmethod
     def is_callable_typehint(typehint):
         typehint = typehint_from_action(typehint)
@@ -917,7 +931,9 @@
     if get_typehint_origin(typehint) in literal_types:
         return typehint  # the args of a Literal are values, not types
     args = getattr(typehint, "__args__", None)
-    if not args:
+    # only a tuple, since e.g. types.UnionType and types.GenericAlias have __args__ as a
+    # class level slot descriptor, which is truthy but not the subtypes of an instance
+    if not isinstance(args, tuple) or not args:
         return typehint
     new_args = tuple(replace_unresolved_forward_refs(a) for a in args)
     if new_args == args:
@@ -991,6 +1007,85 @@
     return required_keys
 
 
+def get_typed_dict_key_type(annotation):
+    # Required and NotRequired only change the requiredness of a key, not its type
+    if get_typehint_origin(annotation) in not_required_required_types:
+        return annotation.__args__[0]
+    return annotation
+
+
+def is_typed_dict_subtype(subtype, typed_dict, logger=None) -> bool:
+    # TypedDicts don't support issubclass, so as specified in PEP 589 the check is done
+    # structurally, i.e. the subtype must have all keys of the typed dict, with the same
+    # types and requiredness.
+    if type(subtype) not in typed_dict_meta_types:
+        return False
+    if subtype is typed_dict:
+        return True
+    annotations = get_typed_dict_annotations(typed_dict, logger)
+    sub_annotations = get_typed_dict_annotations(subtype, logger)
+    for key, annotation in annotations.items():
+        if key not in sub_annotations:
+            return False
+        if get_typed_dict_key_type(sub_annotations[key]) != get_typed_dict_key_type(annotation):
+            return False
+    required_keys = get_typed_dict_required_keys(typed_dict, annotations)
+    sub_required_keys = get_typed_dict_required_keys(subtype, sub_annotations)
+    return required_keys == sub_required_keys & annotations.keys()
+
+
+def is_importable_module_path(val) -> bool:
+    """Whether a value is the import path of a module, checked without importing it.
+
+    Only the parent packages of the module get imported, which is unavoidable
+    since they are the ones that know how to find their submodules.
+    """
+    if not isinstance(val, str) or not all(p.isidentifier() for p in val.split(".")):
+        return False
+    if val in sys.modules:
+        return True
+    try:
+        return find_spec(val) is not None
+    except (ImportError, AttributeError, TypeError, ValueError):
+        return False
+
+
+type_expression_types = {UnionType: "UnionType", GenericAlias: "GenericAlias"}
+
+
+def resolve_type_expression_node(node):
+    """Returns the type that an ast node of a type expression represents."""
+    if isinstance(node, ast.Constant):
+        return NoneType if node.value is None else node.value
+    if isinstance(node, ast.Name):
+        for namespace in (builtins, typing):
+            if hasattr(namespace, node.id):
+                return getattr(namespace, node.id)
+        raise ValueError(f"Not a builtin or typing name: {node.id}")
+    if isinstance(node, ast.Attribute):
+        return import_object(ast.unparse(node))
+    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
+        return resolve_type_expression_node(node.left) | resolve_type_expression_node(node.right)
+    if isinstance(node, ast.Subscript):
+        return resolve_type_expression_node(node.value)[resolve_type_expression_node(node.slice)]
+    if isinstance(node, ast.Tuple):
+        return tuple(resolve_type_expression_node(e) for e in node.elts)
+    if isinstance(node, ast.List):
+        return [resolve_type_expression_node(e) for e in node.elts]
+    raise ValueError(f"Unsupported type expression: {ast.unparse(node)}")
+
+
+def str_to_type_expression(val: str):
+    """Returns the type that a string type expression represents, e.g. ``"int | str"``.
+
+    The expression is resolved from its ast instead of being evaluated, such that only
+    names, dot import paths, unions and subscripts are accepted, i.e. no arbitrary code.
+    """
+    if not isinstance(val, str):
+        raise ValueError(f"Expected a string, got {type(val)}")
+    return resolve_type_expression_node(ast.parse(val, mode="eval").body)
+
+
 def adapt_typehints(
     val,
     typehint,
@@ -1004,7 +1099,8 @@
     default=None,
     logger=None,
 ):
-    if type(val) in {str, bool, int, float} and val == default:
+    # A module import path equal to the default still needs to be imported on instantiation
+    if type(val) in {str, bool, int, float} and val == default and not (instantiate_classes and typehint is ModuleType):
         return val
 
     adapt_kwargs = {
@@ -1088,11 +1184,41 @@
         elif not serialize and not isinstance(val, type):
             path = val
             val = import_object(val)
-            if (typehint in {Type, type} and not isinstance(val, type)) or (
-                typehint not in {Type, type} and not is_subclass(val, subtypehints[0])
-            ):
+            if typehint in {Type, type}:
+                valid = isinstance(val, type)
+            elif type(subtypehints[0]) in typed_dict_meta_types:
+                valid = is_typed_dict_subtype(val, subtypehints[0], logger)
+            else:
+                valid = is_subclass(val, subtypehints[0])
+            if not valid:
                 raise_unexpected_value(f"Expected an import path corresponding to a {typehint}", path)
 
+    # Module
+    elif typehint is ModuleType:
+        if serialize:
+            if isinstance(val, ModuleType):
+                val = val.__name__
+        elif not isinstance(val, ModuleType):
+            if not is_importable_module_path(val):
+                raise_unexpected_value("Expected an import path corresponding to a module", val)
+            if instantiate_classes:
+                val = import_module(val)
+
+    # UnionType and GenericAlias
+    elif typehint in type_expression_types:
+        if serialize:
+            if isinstance(val, typehint):
+                val = str(val)
+        elif not isinstance(val, typehint):
+            expected = f"Expected a string with a {type_expression_types[typehint]} type expression"
+            try:
+                type_expression = str_to_type_expression(val)
+            except Exception as ex:
+                raise_unexpected_value(expected, val, ex)
+            if not isinstance(type_expression, typehint):
+                raise_unexpected_value(expected, val)
+            val = type_expression
+
     # Union
     elif typehint_origin == Union:
         vals = []
@@ -2020,6 +2146,13 @@
     )
 
 
+def is_module_type(annotation):
+    annotation = get_unaliased_type(annotation)
+    return annotation is ModuleType or (
+        get_typehint_origin(annotation) == Union and any(a is ModuleType for a in annotation.__args__)
+    )
+
+
 def is_callable_type(annotation):
     def is_callable(a):
         return (get_typehint_origin(a) or a) in callable_origin_types or a in callable_origin_types
@@ -2040,6 +2173,10 @@
 
 
 def type_to_str(obj):
+    if obj is ModuleType:
+        return "ModuleType"
+    if obj in type_expression_types:
+        return type_expression_types[obj]
     if obj in {bool, tuple} or is_subclass(obj, (int, float, str, Path, Enum)):
         return obj.__name__
     return strip_module_names(str(obj)).replace("NoneType", "null")

jsonargparse_tests/test_link_arguments.py

@@ -2,6 +2,7 @@
 
 import json
 from dataclasses import dataclass
+from types import ModuleType
 from typing import Any, Callable, List, Mapping, Optional, Union
 
 import pytest
@@ -1018,6 +1019,45 @@ def test_on_instantiate_target_entire_dataclass(parser, tmp_cwd):
     assert "--container.dep" not in help_str
 
 
+class ModuleUser:
+    def __init__(self, mod: ModuleType, num: int = 1):
+        self.mod = mod
+        self.num = num
+
+
+def test_on_instantiate_source_module_type(parser):
+    parser.add_argument("--mod", type=ModuleType)
+    parser.add_class_arguments(ModuleUser, "user")
+    parser.link_arguments("mod", "user.mod", apply_on="instantiate")
+
+    cfg = parser.parse_args(["--mod=json"])
+    assert cfg.mod == "json"
+    init = parser.instantiate(cfg)
+    assert init.user.mod is json
+
+
+def test_on_instantiate_source_module_type_compute_fn(parser):
+    parser.add_argument("--mod", type=ModuleType)
+    parser.add_class_arguments(ModuleUser, "user")
+    parser.link_arguments("mod", "user.mod", compute_fn=lambda m: m.decoder, apply_on="instantiate")
+
+    init = parser.instantiate(parser.parse_args(["--mod=json"]))
+    assert init.user.mod is json.decoder
+
+
+def test_on_instantiate_source_module_type_target_subclass(parser):
+    # the module argument is added after the target, so only the instantiation
+    # order given by the link makes the module be imported before it is used
+    parser.add_subclass_arguments(ModuleUser, "user")
+    parser.add_argument("--mod", type=ModuleType)
+    parser.link_arguments("mod", "user.init_args.mod", apply_on="instantiate")
+
+    cfg = parser.parse_args([f"--user={__name__}.ModuleUser", "--mod=json"])
+    init = parser.instantiate(cfg)
+    assert isinstance(init.user, ModuleUser)
+    assert init.user.mod is json
+
+
 # link creation failures
 
 

jsonargparse_tests/test_postponed_annotations.py

@@ -8,7 +8,7 @@
 import typing
 from collections.abc import Callable
 from textwrap import dedent
-from types import SimpleNamespace
+from types import GenericAlias, SimpleNamespace, UnionType
 from typing import TYPE_CHECKING, Dict, ForwardRef, List, Optional, Tuple, Type, TypedDict, Union
 from unittest.mock import patch
 
@@ -428,6 +428,14 @@ def __repr__(self):
         return "Unrebuildable[MisspelledType]"
 
 
+def test_types_with_args_slot_descriptor_unchanged():
+    # types.UnionType and types.GenericAlias have __args__ as a class level slot
+    # descriptor, which is truthy but not the tuple of subtypes of an instance
+    assert replace_unresolved_forward_refs(UnionType) is UnionType
+    assert replace_unresolved_forward_refs(GenericAlias) is GenericAlias
+    assert replace_unresolved_forward_refs(Union[type, UnionType]) == Union[type, UnionType]
+
+
 def test_unresolvable_subtype_not_rebuildable():
     # failing to be rebuilt, the entire type hint becomes unresolved instead of an error
     unresolved = replace_unresolved_forward_refs(UnrebuildableTypehint())