Fix TypedDict requiredness and TYPE_CHECKING bugs

CHANGELOG.rst

@@ -36,6 +36,17 @@ Fixed
 - Fix subclass instance defaults not being converted to dict format when the
   parameter is keyword-only (declared after ``*``) (`#933
   <https://github.com/mauvilsa/jsonargparse/pull/933>`__).
+- Signature ``**kwargs: Unpack[TypedDict]`` where the ``TypedDict`` has
+  ``total=False`` incorrectly made the non-``Required`` keys required (`#934
+  <https://github.com/mauvilsa/jsonargparse/pull/934>`__).
+- ``TypedDict`` requiredness of keys not correctly resolved when using
+  postponed annotations (``from __future__ import annotations``) together with
+  ``Required``/``NotRequired`` or totality-flipping inheritance, affecting both
+  ``Unpack`` and ``type=`` cases (`#934
+  <https://github.com/mauvilsa/jsonargparse/pull/934>`__).
+- ``TypedDict`` keys annotated with types only imported in ``TYPE_CHECKING``
+  blocks not being resolved (`#934
+  <https://github.com/mauvilsa/jsonargparse/pull/934>`__).
 
 Changed
 ^^^^^^^

jsonargparse/_parameter_resolvers.py

@@ -316,16 +316,29 @@ def replace_type_vars(annotation):
             param.annotation = replace_type_vars(param.annotation)
 
 
-def unpack_typed_dict_kwargs(params: ParamList, kwargs_idx: int) -> int:
+def unpack_typed_dict_kwargs(params: ParamList, kwargs_idx: int, logger=None) -> int:
+    from ._typehints import (
+        NotRequired,
+        get_typed_dict_annotations,
+        get_typed_dict_required_keys,
+        not_required_types,
+    )
+
     kwargs = params[kwargs_idx]
     annotation = kwargs.annotation
     if is_unpack_typehint(annotation):
         params.pop(kwargs_idx)
         annotation_args: tuple = getattr(annotation, "__args__", ())
         assert len(annotation_args) == 1, "Unpack requires a single type argument"
-        dict_annotations = annotation_args[0].__annotations__
+        typed_dict = annotation_args[0]
+        dict_annotations = get_typed_dict_annotations(typed_dict, logger)
+        required_keys = get_typed_dict_required_keys(typed_dict, dict_annotations)
         new_params = []
         for nm, annot in dict_annotations.items():
+            if nm not in required_keys and get_typehint_origin(annot) not in not_required_types:
+                # Mark optional keys (e.g. from total=False) as NotRequired so that they
+                # are added as non-required arguments.
+                annot = NotRequired[annot]
             new_params.append(
                 ParamData(
                     name=nm,
@@ -878,7 +891,7 @@ def get_parameters(self) -> ParamList:
         )
         self.replace_param_default_subclass_specs(params)
         if kwargs_idx >= 0:
-            kwargs_idx = unpack_typed_dict_kwargs(params, kwargs_idx)
+            kwargs_idx = unpack_typed_dict_kwargs(params, kwargs_idx, self.logger)
         if args_idx >= 0 or kwargs_idx >= 0:
             self.doc_params = doc_params
             with mro_context(self.parent):

jsonargparse/_typehints.py

@@ -99,6 +99,7 @@
 Required = typing_extensions_import("Required")
 _TypedDictMeta = typing_extensions_import("_TypedDictMeta")
 Unpack = typing_extensions_import("Unpack")
+get_type_hints = typing_extensions_import("get_type_hints")
 
 
 def _capture_typing_extension_shadows(name: str, *collections) -> None:
@@ -781,15 +782,46 @@ def raise_union_unexpected_value(subtypes, val: Any, exceptions: list[Exception]
     ) from exceptions[0]
 
 
-def resolve_forward_ref(ref):
+def resolve_forward_ref(ref, global_vars=None):
     if not isinstance(ref, ForwardRef) or not ref.__forward_module__:
         return ref
 
     aliases = __builtins__.copy()
     aliases.update(vars(import_module(ref.__forward_module__)))
+    if global_vars:
+        aliases.update(global_vars)
     return aliases.get(ref.__forward_arg__, ref)
 
 
+def get_typed_dict_annotations(typed_dict, logger=None) -> dict:
+    from ._postponed_annotations import get_global_vars
+
+    # Includes the names from TYPE_CHECKING blocks, given as localns so that each base
+    # keeps resolving with the globals of the module in which it was defined.
+    global_vars = get_global_vars(typed_dict, logger)
+    try:
+        # Resolves forward references (e.g. from "from __future__ import annotations") and
+        # gathers inherited keys, while include_extras keeps the Required/NotRequired wrappers.
+        return get_type_hints(typed_dict, None, global_vars, include_extras=True)
+    except Exception as ex:
+        if logger:
+            logger.debug(f"Failed to resolve the annotations of {typed_dict}", exc_info=ex)
+    # A single key failing (e.g. a missing import or a typo) makes the resolution of the
+    # entire TypedDict fail. Thus, resolve one by one to keep the keys that do work.
+    return {k: resolve_forward_ref(v, global_vars) for k, v in typed_dict.__annotations__.items()}
+
+
+def get_typed_dict_required_keys(typed_dict, annotations: dict) -> set:
+    # The totality of each class (including inheritance) is reflected in __required_keys__,
+    # even when the annotations are postponed. Required and NotRequired instead may not be
+    # reflected there (e.g. below Python 3.11 or with postponed annotations), so they are
+    # adjusted based on the resolved annotations.
+    required_keys = set(getattr(typed_dict, "__required_keys__", set(annotations)))
+    required_keys.update({k for k, v in annotations.items() if get_typehint_origin(v) in required_types})
+    required_keys.difference_update({k for k, v in annotations.items() if get_typehint_origin(v) in not_required_types})
+    return required_keys
+
+
 def adapt_typehints(
     val,
     typehint,
@@ -1009,27 +1041,16 @@ def adapt_typehints(
                         kwargs["prev_val"] = None
                 val[k] = adapt_typehints(v, subtypehints[1], **kwargs)
         if type(typehint) in typed_dict_meta_types:
-            if hasattr(typehint, "__required_keys__"):
-                required_keys = set(typehint.__required_keys__)
-                # The standard library TypedDict below Python 3.11 does not store runtime
-                # information about optional and required keys when using Required or NotRequired.
-                # Thus, capture explicitly Required keys
-                required_keys.update(
-                    {k for k, v in typehint.__annotations__.items() if get_typehint_origin(v) in required_types}
-                )
-                # And remove explicitly NotRequired keys
-                required_keys.difference_update(
-                    {k for k, v in typehint.__annotations__.items() if get_typehint_origin(v) in not_required_types}
-                )
+            dict_annotations = get_typed_dict_annotations(typehint, logger)
+            required_keys = get_typed_dict_required_keys(typehint, dict_annotations)
             missing_keys = required_keys - val.keys()
             if missing_keys:
                 raise_unexpected_value(f"Missing required keys: {missing_keys}", val)
-            extra_keys = val.keys() - typehint.__annotations__.keys()
+            extra_keys = val.keys() - dict_annotations.keys()
             if extra_keys:
                 raise_unexpected_value(f"Unexpected keys: {extra_keys}", val)
             for k, v in val.items():
-                subtypehint = resolve_forward_ref(typehint.__annotations__[k])
-                val[k] = adapt_typehints(v, subtypehint, **adapt_kwargs)
+                val[k] = adapt_typehints(v, dict_annotations[k], **adapt_kwargs)
         if typehint_origin is MappingProxyType and not serialize:
             val = MappingProxyType(val)
         elif typehint_origin is OrderedDict:

jsonargparse_tests/test_postponed_annotations.py

@@ -6,7 +6,7 @@
 import sys
 from textwrap import dedent
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Dict, ForwardRef, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Dict, ForwardRef, List, Optional, Tuple, Type, TypedDict, Union
 from unittest.mock import patch
 
 import pytest
@@ -26,6 +26,7 @@
     get_types,
     type_requires_eval,
 )
+from jsonargparse._typehints import Unpack, get_typed_dict_annotations
 from jsonargparse.typing import Path_drw
 from jsonargparse_tests.conftest import capture_logs, source_unavailable
 from jsonargparse_tests.test_dataclasses import DifferentModuleBaseData
@@ -149,6 +150,7 @@ def test_type_checking_visitor_failure(logger):
 
 if TYPE_CHECKING:  # pragma: no cover
     import xml.dom
+    from decimal import Decimal
 
     class TypeCheckingClass1:
         pass
@@ -240,6 +242,95 @@ def test_get_types_type_checking_tuple():
     assert str(types["p1"]) == f"{tpl}[{__name__}.TypeCheckingClass1, {__name__}.TypeCheckingClass2]"
 
 
+# Types only imported in a TYPE_CHECKING block are resolved for TypedDict keys, the same
+# as for regular parameters of a signature.
+
+
+class TypeCheckingTypedDict(TypedDict, total=False):
+    num: int
+    amount: Decimal
+    only_typing: TypeCheckingClass1
+
+
+class TypeCheckingTypedDictClass:
+    def __init__(self, **kwargs: Unpack[TypeCheckingTypedDict]) -> None:
+        self.kwargs = kwargs  # pragma: no cover
+
+
+def test_get_typed_dict_annotations_type_checking():
+    annotations = get_typed_dict_annotations(TypeCheckingTypedDict)
+    assert annotations["num"] is int
+    assert annotations["amount"] is __import__("decimal").Decimal
+    assert annotations["only_typing"].__name__ == "TypeCheckingClass1"
+
+
+def test_typed_dict_type_checking_type(parser):
+    parser.add_argument("--opts", type=TypeCheckingTypedDict)
+    cfg = parser.parse_args(['--opts={"num": 1, "amount": "1.5"}'])
+    assert cfg.opts == {"num": 1, "amount": __import__("decimal").Decimal("1.5")}
+
+
+@pytest.mark.skipif(not Unpack, reason="Unpack introduced in python 3.11 or backported in typing_extensions")
+def test_typed_dict_type_checking_unpack(parser):
+    added = parser.add_class_arguments(TypeCheckingTypedDictClass, "cls")
+    assert added == ["cls.num", "cls.amount", "cls.only_typing"]
+    cfg = parser.parse_args(["--cls.num=1", "--cls.amount=1.5"])
+    assert cfg.cls == Namespace(num=1, amount=__import__("decimal").Decimal("1.5"))
+
+
+# When an annotation can't be resolved, e.g. a missing import or a typo, get_type_hints
+# fails for the entire TypedDict. Thus, the annotations are resolved one by one so that
+# the keys that do resolve remain usable.
+
+
+class UnresolvableTypedDict(TypedDict, total=False):
+    num: int
+    typo: MisspelledType  # type: ignore[name-defined]  # noqa: F821
+
+
+func_unresolvable_typed_dict = TypedDict(
+    "func_unresolvable_typed_dict",
+    {"num": int, "ref": "Path_drw", "typo": "MisspelledType"},  # type: ignore[name-defined]  # noqa: F821
+    total=False,
+)
+
+
+def assert_unresolved_forward_ref(annotation, name):
+    assert isinstance(annotation, ForwardRef)
+    assert annotation.__forward_arg__ == name
+
+
+def test_get_typed_dict_annotations_unresolvable_key():
+    annotations = get_typed_dict_annotations(UnresolvableTypedDict)
+    assert annotations["num"] is int
+    assert_unresolved_forward_ref(annotations["typo"], "MisspelledType")
+
+
+def test_get_typed_dict_annotations_unresolvable_key_functional():
+    annotations = get_typed_dict_annotations(func_unresolvable_typed_dict)
+    assert annotations["num"] is int  # not a forward ref, used as is
+    assert annotations["ref"] is Path_drw  # forward ref resolved from the module
+    assert_unresolved_forward_ref(annotations["typo"], "MisspelledType")
+
+
+def test_typed_dict_unresolvable_key_type(parser):
+    parser.add_argument("--opts", type=UnresolvableTypedDict)
+    cfg = parser.parse_args(['--opts={"num": 1}'])
+    assert cfg.opts == {"num": 1}
+
+
+class UnresolvableTypedDictClass:
+    def __init__(self, **kwargs: Unpack[UnresolvableTypedDict]) -> None:
+        self.kwargs = kwargs  # pragma: no cover
+
+
+@pytest.mark.skipif(not Unpack, reason="Unpack introduced in python 3.11 or backported in typing_extensions")
+def test_typed_dict_unresolvable_key_unpack(parser):
+    added = parser.add_class_arguments(UnresolvableTypedDictClass, "cls")
+    assert added == ["cls.num"]  # the unresolvable key is skipped
+    assert parser.parse_args(["--cls.num=1"]).cls == Namespace(num=1)
+
+
 def function_type_checking_type(p1: Type["TypeCheckingClass2"]):
     return p1  # pragma: no cover