From b985bb847cd0498c254e1228d57b5b3815b78415 Mon Sep 17 00:00:00 2001 From: Mauricio Villegas <5780272+mauvilsa@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:56:22 +0200 Subject: [PATCH 1/3] Add support for type[SomeTypedDict] --- CHANGELOG.rst | 4 ++ DOCUMENTATION.rst | 8 ++- jsonargparse/_typehints.py | 37 +++++++++++-- jsonargparse_tests/test_typehints.py | 77 ++++++++++++++++++++++++++++ 4 files changed, 122 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b1c8b615..ed9b0780 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -29,6 +29,10 @@ Added - ``shtab`` completion scripts now include file and directory completions for arguments typed as pydantic's ``FilePath`` and ``DirectoryPath`` (`#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 (`#??? + `__). Fixed ^^^^^ diff --git a/DOCUMENTATION.rst b/DOCUMENTATION.rst index fbcbb8ea..3a230d9a 100644 --- a/DOCUMENTATION.rst +++ b/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 `__. - 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 `__, + 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`` diff --git a/jsonargparse/_typehints.py b/jsonargparse/_typehints.py index dec3b93d..36e7cbde 100644 --- a/jsonargparse/_typehints.py +++ b/jsonargparse/_typehints.py @@ -991,6 +991,33 @@ def get_typed_dict_required_keys(typed_dict, annotations: dict) -> set: 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 adapt_typehints( val, typehint, @@ -1088,9 +1115,13 @@ def adapt_typehints( 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) # Union diff --git a/jsonargparse_tests/test_typehints.py b/jsonargparse_tests/test_typehints.py index b36af237..75c0f633 100644 --- a/jsonargparse_tests/test_typehints.py +++ b/jsonargparse_tests/test_typehints.py @@ -47,6 +47,7 @@ get_all_subclass_paths, get_subclass_types, is_optional, + is_typed_dict_subtype, type_to_str, ) from jsonargparse._util import get_import_path @@ -755,6 +756,82 @@ def test_typeddict_with_required_arg(parser): ctx.match("Expected a ") +# type[TypedDict] tests. TypedDicts don't support issubclass, so the check is structural. + + +class StateDict(TypedDict): + messages: list + + +class SubStateDict(StateDict): + extra: int + + +class SameKeysDict(TypedDict): + messages: list + + +class DifferentTypeDict(TypedDict): + messages: dict + + +class MissingKeyDict(TypedDict): + extra: int + + +class NotTotalStateDict(TypedDict, total=False): + messages: list + + +def test_type_typeddict_accepts_self_and_subclass(parser): + parser.add_argument("--cls", type=Type[StateDict]) + assert parser.parse_args([f"--cls={__name__}.StateDict"]).cls is StateDict + assert parser.parse_args([f"--cls={__name__}.SubStateDict"]).cls is SubStateDict + assert json_or_yaml_load(parser.dump(parser.parse_args([f"--cls={__name__}.SubStateDict"]))) == { + "cls": f"{__name__}.SubStateDict" + } + + +def test_type_typeddict_accepts_structurally_equivalent(parser): + parser.add_argument("--cls", type=Type[StateDict]) + assert parser.parse_args([f"--cls={__name__}.SameKeysDict"]).cls is SameKeysDict + + +def test_type_typeddict_rejects_incompatible(parser): + parser.add_argument("--cls", type=Type[StateDict]) + for name in ["DifferentTypeDict", "MissingKeyDict", "NotTotalStateDict"]: + with pytest.raises(ArgumentError) as ctx: + parser.parse_args([f"--cls={__name__}.{name}"]) + ctx.match("Expected an import path corresponding to a") + with pytest.raises(ArgumentError) as ctx: + parser.parse_args(["--cls=uuid.UUID"]) + ctx.match("Expected an import path corresponding to a") + + +def test_type_typeddict_optional(parser): + parser.add_argument("--cls", type=Optional[Type[StateDict]], default=None) + assert parser.parse_args([]).cls is None + assert parser.parse_args(["--cls=null"]).cls is None + assert parser.parse_args([f"--cls={__name__}.SubStateDict"]).cls is SubStateDict + + +@pytest.mark.skipif(not NotRequired, reason="NotRequired introduced in python 3.11 or backported in typing_extensions") +def test_is_typed_dict_subtype_not_required_key(): + base = TypedDict("BaseNotRequiredDict", {"a": NotRequired[int]}) + not_total = TypedDict("NotTotalDict", {"a": int}, total=False) + total = TypedDict("TotalDict", {"a": int}) + assert is_typed_dict_subtype(not_total, base) + assert not is_typed_dict_subtype(total, base) + + +def test_type_typeddict_help(parser): + parser.add_argument("--cls", type=Optional[Type[StateDict]], default=None) + help_str = get_parser_help(parser) + assert "--cls CLS" in help_str + assert "StateDict" in help_str + assert "default: null" in help_str + + # Required/NotRequired as the type of an argument. The wrapper must agree with the # requiredness of the argument and is removed so that it is not shown in the help. From f7816e86e0aac9a0823d4d346679683e014455ce Mon Sep 17 00:00:00 2001 From: Mauricio Villegas <5780272+mauvilsa@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:19:48 +0200 Subject: [PATCH 2/3] Add support for ModuleType --- CHANGELOG.rst | 4 + DOCUMENTATION.rst | 3 + jsonargparse/_link_arguments.py | 2 +- jsonargparse/_typehints.py | 50 ++++++++- jsonargparse_tests/test_link_arguments.py | 40 ++++++++ jsonargparse_tests/test_typehints.py | 120 +++++++++++++++++++++- 6 files changed, 215 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ed9b0780..a20b4180 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -33,6 +33,10 @@ Added is structurally compatible, i.e. it has all the keys of the expected ``TypedDict``, with the same types and requiredness (`#??? `__). +- 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 (`#??? + `__). Fixed ^^^^^ diff --git a/DOCUMENTATION.rst b/DOCUMENTATION.rst index 3a230d9a..8f9c559c 100644 --- a/DOCUMENTATION.rst +++ b/DOCUMENTATION.rst @@ -597,6 +597,9 @@ 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. + - ``TypeAliasType`` is supported with values parsed as the aliased type and the alias shown as the argument type in help. diff --git a/jsonargparse/_link_arguments.py b/jsonargparse/_link_arguments.py index b33c2499..328621e6 100644 --- a/jsonargparse/_link_arguments.py +++ b/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: diff --git a/jsonargparse/_typehints.py b/jsonargparse/_typehints.py index 36e7cbde..0bc6ca55 100644 --- a/jsonargparse/_typehints.py +++ b/jsonargparse/_typehints.py @@ -12,7 +12,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, MappingProxyType, ModuleType from typing import ( Any, Callable, @@ -151,6 +152,7 @@ def _capture_typing_extension_shadows(name: str, *collections) -> None: OrderedDict, Callable, abc.Callable, + ModuleType, NotRequired, Required, Unpack, @@ -357,6 +359,8 @@ def normalize_default(self, default): 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 +465,11 @@ def is_mapping_typehint(typehint): 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) @@ -1018,6 +1027,22 @@ def is_typed_dict_subtype(subtype, typed_dict, logger=None) -> bool: 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 + + def adapt_typehints( val, typehint, @@ -1031,7 +1056,8 @@ def adapt_typehints( 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 = { @@ -1124,6 +1150,17 @@ def adapt_typehints( 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) + # Union elif typehint_origin == Union: vals = [] @@ -2051,6 +2088,13 @@ def is_enum_type(annotation): ) +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 @@ -2071,6 +2115,8 @@ def strip_module_names(string: str) -> str: def type_to_str(obj): + if obj is ModuleType: + return "ModuleType" 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") diff --git a/jsonargparse_tests/test_link_arguments.py b/jsonargparse_tests/test_link_arguments.py index 23d3652e..8be0713a 100644 --- a/jsonargparse_tests/test_link_arguments.py +++ b/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 diff --git a/jsonargparse_tests/test_typehints.py b/jsonargparse_tests/test_typehints.py index 75c0f633..ca4c2464 100644 --- a/jsonargparse_tests/test_typehints.py +++ b/jsonargparse_tests/test_typehints.py @@ -12,7 +12,7 @@ from enum import Enum from pathlib import Path from textwrap import dedent -from types import MappingProxyType +from types import MappingProxyType, ModuleType from typing import ( Any, Callable, @@ -832,6 +832,124 @@ def test_type_typeddict_help(parser): assert "default: null" in help_str +# ModuleType tests. The value is the import path of a module, which is only +# imported when instantiate_classes is run. + + +@pytest.fixture +def unimported_module(tmp_path, monkeypatch): + name = "jsonargparse_tests_unimported_module" + (tmp_path / f"{name}.py").write_text("value = 3\n") + monkeypatch.syspath_prepend(str(tmp_path)) + importlib.invalidate_caches() + assert name not in sys.modules + yield name + sys.modules.pop(name, None) + + +def test_module_type_parse_keeps_import_path(parser): + parser.add_argument("--mod", type=ModuleType) + cfg = parser.parse_args(["--mod=json"]) + assert cfg.mod == "json" + + +def test_module_type_parse_submodule(parser): + parser.add_argument("--mod", type=ModuleType) + cfg = parser.parse_args(["--mod=json.decoder"]) + assert cfg.mod == "json.decoder" + init = parser.instantiate(cfg) + assert init.mod is json.decoder + + +def test_module_type_not_imported_on_parse(parser, unimported_module): + parser.add_argument("--mod", type=ModuleType) + cfg = parser.parse_args([f"--mod={unimported_module}"]) + assert cfg.mod == unimported_module + assert unimported_module not in sys.modules + + +def test_module_type_instantiate_imports_module(parser, unimported_module): + parser.add_argument("--mod", type=ModuleType) + cfg = parser.parse_args([f"--mod={unimported_module}"]) + init = parser.instantiate(cfg) + assert isinstance(init.mod, ModuleType) + assert init.mod.value == 3 + assert unimported_module in sys.modules + + +@pytest.mark.parametrize("value", ["not_a_module", "uuid.UUID", "json.not_a_submodule", "not.a.module", "", "1json"]) +def test_module_type_invalid_import_path(parser, value): + parser.add_argument("--mod", type=ModuleType) + with pytest.raises(ArgumentError) as ctx: + parser.parse_args([f"--mod={value}"]) + ctx.match("Expected an import path corresponding to a module") + + +def test_module_type_optional(parser): + parser.add_argument("--mod", type=Optional[ModuleType], default=None) + assert parser.parse_args([]).mod is None + assert parser.parse_args(["--mod=null"]).mod is None + cfg = parser.parse_args(["--mod=json"]) + assert cfg.mod == "json" + assert parser.instantiate(cfg).mod is json + + +def test_module_type_default_module_object(parser): + parser.add_argument("--mod", type=ModuleType, default=json) + cfg = parser.parse_args([]) + assert cfg.mod == "json" + assert parser.instantiate(cfg).mod is json + + +def test_module_type_list(parser): + parser.add_argument("--mods", type=List[ModuleType], default=[]) + cfg = parser.parse_args(['--mods=["json", "uuid"]']) + assert cfg.mods == ["json", "uuid"] + assert parser.instantiate(cfg).mods == [json, uuid] + + +def test_module_type_dump(parser): + parser.add_argument("--mod", type=ModuleType) + cfg = parser.parse_args(["--mod=json"]) + assert json_or_yaml_load(parser.dump(cfg)) == {"mod": "json"} + + +def test_module_type_dump_module_object(parser): + parser.add_argument("--mod", type=ModuleType) + cfg = parser.parse_args(["--mod=json"]) + cfg.mod = json + assert json_or_yaml_load(parser.dump(cfg)) == {"mod": "json"} + + +def test_module_type_help(parser): + parser.add_argument("--mod", type=ModuleType, help="Module to use.") + help_str = get_parser_help(parser) + assert "--mod MOD" in help_str + assert "Module to use. (type: ModuleType, default: null)" in help_str + + +class WithModule: + def __init__(self, mod: ModuleType, num: int = 1): + self.mod = mod + self.num = num + + +def test_module_type_class_group_instantiate(parser): + parser.add_class_arguments(WithModule, "cls") + cfg = parser.parse_args(["--cls.mod=json"]) + assert cfg.cls.mod == "json" + init = parser.instantiate(cfg) + assert init.cls.mod is json + + +def test_module_type_subclass_init_arg_instantiate(parser): + parser.add_argument("--cls", type=WithModule) + cfg = parser.parse_args([f"--cls={__name__}.WithModule", "--cls.mod=json"]) + assert cfg.cls.init_args.mod == "json" + init = parser.instantiate(cfg) + assert init.cls.mod is json + + # Required/NotRequired as the type of an argument. The wrapper must agree with the # requiredness of the argument and is removed so that it is not shown in the help. From 3ca39dc397c8388f1d5bdf493430c0e41cd355d7 Mon Sep 17 00:00:00 2001 From: Mauricio Villegas <5780272+mauvilsa@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:10:45 +0200 Subject: [PATCH 3/3] Add support for UnionType and GenericAlias --- CHANGELOG.rst | 14 ++- DOCUMENTATION.rst | 6 + jsonargparse/_typehints.py | 64 ++++++++++- .../test_postponed_annotations.py | 10 +- jsonargparse_tests/test_typehints.py | 105 +++++++++++++++++- 5 files changed, 190 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a20b4180..13dd1380 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -31,12 +31,18 @@ Added `__). - 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 (`#??? - `__). + ``TypedDict``, with the same types and requiredness (`#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 (`#??? - `__). + is replaced by the imported module object (`#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 + `__). Fixed ^^^^^ diff --git a/DOCUMENTATION.rst b/DOCUMENTATION.rst index 8f9c559c..49e2eb6f 100644 --- a/DOCUMENTATION.rst +++ b/DOCUMENTATION.rst @@ -600,6 +600,12 @@ Some notes about this support are: - ``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. diff --git a/jsonargparse/_typehints.py b/jsonargparse/_typehints.py index 0bc6ca55..ab940944 100644 --- a/jsonargparse/_typehints.py +++ b/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 @@ -13,7 +16,7 @@ from functools import partial from importlib import import_module from importlib.util import find_spec -from types import FunctionType, MappingProxyType, ModuleType +from types import FunctionType, GenericAlias, MappingProxyType, ModuleType, UnionType from typing import ( Any, Callable, @@ -153,6 +156,8 @@ def _capture_typing_extension_shadows(name: str, *collections) -> None: Callable, abc.Callable, ModuleType, + UnionType, + GenericAlias, NotRequired, Required, Unpack, @@ -926,7 +931,9 @@ def replace_unresolved_forward_refs(typehint): 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: @@ -1043,6 +1050,42 @@ def is_importable_module_path(val) -> bool: 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, @@ -1161,6 +1204,21 @@ def adapt_typehints( 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 = [] @@ -2117,6 +2175,8 @@ def strip_module_names(string: str) -> str: 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") diff --git a/jsonargparse_tests/test_postponed_annotations.py b/jsonargparse_tests/test_postponed_annotations.py index cd9a3ca4..0ac36e8c 100644 --- a/jsonargparse_tests/test_postponed_annotations.py +++ b/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()) diff --git a/jsonargparse_tests/test_typehints.py b/jsonargparse_tests/test_typehints.py index ca4c2464..a7e3f767 100644 --- a/jsonargparse_tests/test_typehints.py +++ b/jsonargparse_tests/test_typehints.py @@ -1,5 +1,6 @@ from __future__ import annotations +import calendar import importlib.util import json import pickle @@ -7,12 +8,12 @@ import sys import time import uuid -from collections import OrderedDict, deque +from collections import OrderedDict, abc, deque from dataclasses import dataclass, field from enum import Enum from pathlib import Path from textwrap import dedent -from types import MappingProxyType, ModuleType +from types import GenericAlias, MappingProxyType, ModuleType, UnionType from typing import ( Any, Callable, @@ -950,6 +951,106 @@ def test_module_type_subclass_init_arg_instantiate(parser): assert init.cls.mod is json +# types.UnionType and types.GenericAlias tests. The value is a string with a type +# expression, e.g. "int | str" and "list[int]". + + +def test_union_type_parse(parser): + parser.add_argument("--type", type=UnionType) + assert parser.parse_args(["--type=int | str"]).type == int | str + assert parser.parse_args(["--type=int|None"]).type == Optional[int] + assert parser.parse_args(["--type=calendar.Calendar | uuid.UUID"]).type == calendar.Calendar | uuid.UUID + + +def test_union_type_parse_subscripted_subtype(parser): + parser.add_argument("--type", type=UnionType) + assert parser.parse_args(["--type=list[int] | str"]).type == list[int] | str + + +@pytest.mark.parametrize("value", ["int", "list[int]", "not_a_type | int", "int |", "1 + 2", "print('x')", ""]) +def test_union_type_invalid(parser, value): + parser.add_argument("--type", type=UnionType) + with pytest.raises(ArgumentError) as ctx: + parser.parse_args([f"--type={value}"]) + ctx.match("Expected a string with a UnionType type expression") + + +def test_union_type_dump(parser): + parser.add_argument("--type", type=UnionType) + cfg = parser.parse_args(["--type=int | str"]) + assert json_or_yaml_load(parser.dump(cfg)) == {"type": "int | str"} + + +def test_union_type_default(parser): + parser.add_argument("--type", type=UnionType, default=int | str) + cfg = parser.parse_args([]) + assert cfg.type == int | str + assert json_or_yaml_load(parser.dump(cfg)) == {"type": "int | str"} + + +def test_union_type_optional(parser): + parser.add_argument("--type", type=Optional[UnionType], default=None) + assert parser.parse_args([]).type is None + assert parser.parse_args(["--type=null"]).type is None + assert parser.parse_args(["--type=int | str"]).type == int | str + + +def test_union_type_help(parser): + parser.add_argument("--type", type=UnionType, help="Type to use.") + help_str = get_parser_help(parser) + assert "--type TYPE" in help_str + assert "Type to use. (type: UnionType, default: null)" in help_str + + +def test_generic_alias_parse(parser): + parser.add_argument("--type", type=GenericAlias) + assert parser.parse_args(["--type=list[int]"]).type == list[int] + assert parser.parse_args(["--type=dict[str, Any]"]).type == dict[str, Any] + assert parser.parse_args(["--type=tuple[int, ...]"]).type == tuple[int, ...] + assert parser.parse_args(["--type=list[calendar.Calendar]"]).type == list[calendar.Calendar] + assert parser.parse_args(["--type=collections.abc.Callable[[int], str]"]).type == abc.Callable[[int], str] + + +@pytest.mark.parametrize("value", ["int", "int | str", "List[int]", "list[not_a_type]", "list[", ""]) +def test_generic_alias_invalid(parser, value): + parser.add_argument("--type", type=GenericAlias) + with pytest.raises(ArgumentError) as ctx: + parser.parse_args([f"--type={value}"]) + ctx.match("Expected a string with a GenericAlias type expression") + + +def test_generic_alias_dump(parser): + parser.add_argument("--type", type=GenericAlias) + cfg = parser.parse_args(["--type=dict[str, int]"]) + assert json_or_yaml_load(parser.dump(cfg)) == {"type": "dict[str, int]"} + + +def test_generic_alias_help(parser): + parser.add_argument("--type", type=GenericAlias, help="Type to use.") + help_str = get_parser_help(parser) + assert "Type to use. (type: GenericAlias, default: null)" in help_str + + +def function_schema(schema: Union[type, UnionType, Dict[str, Any]] = int): + return schema # pragma: no cover + + +def test_type_or_union_type_or_dict_function(parser): + added = parser.add_function_arguments(function_schema, "fn") + assert added == ["fn.schema"] + assert parser.parse_args([]).fn.schema is int + assert parser.parse_args(["--fn.schema=calendar.Calendar"]).fn.schema is calendar.Calendar + assert parser.parse_args(["--fn.schema=int | str"]).fn.schema == int | str + assert parser.parse_args(['--fn.schema={"key": 1}']).fn.schema == {"key": 1} + + +def test_union_type_list(parser): + parser.add_argument("--types", type=List[UnionType], default=[]) + cfg = parser.parse_args(['--types=["int | str", "float | None"]']) + assert cfg.types == [int | str, Optional[float]] + assert json_or_yaml_load(parser.dump(cfg)) == {"types": ["int | str", "float | None"]} + + # Required/NotRequired as the type of an argument. The wrapper must agree with the # requiredness of the argument and is removed so that it is not shown in the help.