Repository: omni-us/jsonargparse (python)

Issue:
`NotRequired` keys of `TypedDict` are incorrectly inferred as required
<!-- If you like this project, please ⭐ star it https://github.com/omni-us/jsonargparse/ -->

## 🐛 Bug report

See the [documentation for TypedDict](https://docs.python.org/3/library/typing.html#typing.TypedDict) for context. 
Note that by default, a `TypedDict` requires all keys to be present, i.e., it is equivalent to being created with `total=True` as a metaclass argument.
One may specify individual keys as not required with the `NotRequired` annotation. 
This annotation is not respected by `jsonargparse`.
Namely, [this check](https://github.com/omni-us/jsonargparse/blob/9919b35f0d554bf90dced859969229e008a4de91/jsonargparse/_typehints.py#L893) in `jsonargparse._typehints` does not consider the possibility of per-field `Required`/`NotRequired` annotations.

### To reproduce

```python

from dataclasses import dataclass, field
from jsonargparse import ActionConfigFile, ArgumentParser, lazy_instance
import tempfile
from typing import TypedDict, NotRequired
import yaml

class TestDict(TypedDict):
    a: int
    b: NotRequired[int]

@dataclass
class TestClass:

    test: TestDict = field(default_factory=lambda: TestDict(a=0))

parser = ArgumentParser(exit_on_error=False)
parser.add_argument(
    "-c",
    "--config",
    action=ActionConfigFile,
    help="Path to a configuration file in json or yaml format.")
parser.add_class_arguments(
    TestClass,
    "test",
    fail_untyped=False,
    instantiate=True,
    sub_configs=True,
    default=lazy_instance(
        TestClass,
    ),
)

config = yaml.safe_dump({
    "test": {
        "test": {
            "a": 1
        }
    }
})

with tempfile.NamedTemporaryFile("w", suffix=".yaml") as f:
    f.write(config)
    f.flush()
    cfg = parser.parse_args(["--config", f"{f.name}"])

```

### Expected behavior

The script should exit without error after successfully parsing the arguments.


### Actual behavior

<details>
<summary> Traceback </summary>

```python
Traceback (most recent call last):
  File ".../lib/python3.11/site-packages/jsonargparse/_typehints.py", line 581, in _check_type
    raise ex
  File ".../lib/python3.11/site-packages/jsonargparse/_typehints.py", line 566, in _check_type
    val = adapt_typehints(val, self._typehint, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_typehints.py", line 913, in adapt_typehints
    raise_unexpected_value(f"Missing required keys: {missing_keys}", val)
  File ".../lib/python3.11/site-packages/jsonargparse/_typehints.py", line 702, in raise_unexpected_value
    raise ValueError(message) from exception
ValueError: Missing required keys: {'b'}. Got value: {'a': 1}

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 605, in parse_string
    cfg = self._load_config_parser_mode(cfg_str, cfg_path, ext_vars, previous_config.get())
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 654, in _load_config_parser_mode
    return self._apply_actions(cfg_dict, prev_cfg=prev_cfg)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 1314, in _apply_actions
    value = self._check_value_key(action, value, action_dest, prev_cfg)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 1367, in _check_value_key
    value = action._check_type_(value, cfg=cfg)  # type: ignore[attr-defined]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_common.py", line 268, in _check_type_
    return self._check_type(value, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_typehints.py", line 591, in _check_type
    raise TypeError(f'Parser key "{self.dest}"{elem}:\n{error}') from ex
TypeError: Parser key "test.test":
  Missing required keys: {'b'}. Got value: {'a': 1}

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 264, in parse_known_args
    namespace, args = self._parse_known_args(args, namespace)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/argparse.py", line 2114, in _parse_known_args
    start_index = consume_optional(start_index)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/argparse.py", line 2054, in consume_optional
    take_action(action, args, option_string)
  File ".../lib/python3.11/argparse.py", line 1978, in take_action
    action(self, namespace, argument_values, option_string)
  File ".../lib/python3.11/site-packages/jsonargparse/_actions.py", line 170, in __call__
    self.apply_config(parser, cfg, self.dest, values)
  File ".../lib/python3.11/site-packages/jsonargparse/_actions.py", line 205, in apply_config
    cfg_file = parser.parse_path(value, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 562, in parse_path
    parsed_cfg = self.parse_string(
                 ^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_deprecated.py", line 123, in patched_parse
    cfg = parse_method(*args, _skip_check=_skip_check, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 621, in parse_string
    self.error(str(ex), ex)
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 1010, in error
    raise argument_error(message) from ex
argparse.ArgumentError: Parser key "test.test":
  Missing required keys: {'b'}. Got value: {'a': 1}

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  ...
  File "...", line 140, in <module>
    cfg = parser.parse_args(["--config", f"{f.name}"])
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_deprecated.py", line 123, in patched_parse
    cfg = parse_method(*args, _skip_check=_skip_check, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 396, in parse_args
    cfg, unk = self.parse_known_args(args=args, namespace=cfg)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 266, in parse_known_args
    self.error(str(ex), ex)
  File ".../lib/python3.11/site-packages/jsonargparse/_core.py", line 1010, in error
    raise argument_error(message) from ex
argparse.ArgumentError: Parser key "test.test":
  Missing required keys: {'b'}. Got value: {'a': 1}
```

</details>

### Environment

<!-- Fill in the list below. -->

- jsonargparse version: 4.32.0
- Python version: 3.11.6
- How jsonargparse was installed: `pip`
- OS: Ubuntu 20.04


Pull request description:
Fix inference of required keys in TypedDicts
<!--
Thank you very much for contributing! If you like this project, please ⭐ star it
https://github.com/omni-us/jsonargparse/
-->

## What does this PR do?

<!--
Concisely describe what this pull request does. If available, include links to
issues web pages where the need for this change has been motivated.
-->

This PR fixes inference of required and missing keys in `TypedDict`s.

Closes #570 

## Before submitting

<!--
Use the following list as tasks to be completed before marking a pull request as
"Ready for review". Fill with an "x" all tasks done. Use "n/a" if you think a
task is not relevant or leave empty when in doubt.
-->

- [x] Did you read the [contributing guideline](https://github.com/omni-us/jsonargparse/blob/main/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)


Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2348ca0..67fbc8a 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -18,6 +18,8 @@ v4.33.0 (2024-09-??)
 Added
 ^^^^^
 - Support for Python 3.13.
+- Support for `NotRequired` and `Required` annotations for `TypedDict` keys (`#571
+  <https://github.com/omni-us/jsonargparse/pull/571>`__)
 
 Fixed
 ^^^^^
diff --git a/DOCUMENTATION.rst b/DOCUMENTATION.rst
index dfb1bde..d8e57b7 100644
--- a/DOCUMENTATION.rst
+++ b/DOCUMENTATION.rst
@@ -401,8 +401,10 @@ Some notes about this support are:
   :ref:`parsing-paths` and :ref:`parsing-urls`.
 
 - ``Dict``, ``Mapping``, ``MutableMapping``, ``MappingProxyType``,
-  ``OrderedDict`` and ``TypedDict`` are supported but only with ``str`` or
-  ``int`` keys. For more details see :ref:`dict-items`.
+  ``OrderedDict``, and ``TypedDict`` are supported but only with ``str`` or
+  ``int`` keys. ``Required`` and ``NotRequired`` are also supported for
+  fine-grained specification of required/optional ``TypedDict`` keys.
+  For more details see :ref:`dict-items`.
 
 - ``Tuple``, ``Set`` and ``MutableSet`` are supported even though they can't be
   represented in json distinguishable from a list. Each ``Tuple`` element
diff --git a/jsonargparse/_typehints.py b/jsonargparse/_typehints.py
index 43dff94..b1a20bc 100644
--- a/jsonargparse/_typehints.py
+++ b/jsonargparse/_typehints.py
@@ -3,6 +3,7 @@
 import inspect
 import os
 import re
+import sys
 from argparse import ArgumentError
 from collections import OrderedDict, abc, defaultdict
 from contextlib import contextmanager, suppress
@@ -88,6 +89,21 @@ __all__ = ["lazy_instance"]
 
 
 Literal = typing_extensions_import("Literal")
+NotRequired = typing_extensions_import("NotRequired")
+Required = typing_extensions_import("Required")
+TypedDict = typing_extensions_import("TypedDict")
+_TypedDictMeta = typing_extensions_import("_TypedDictMeta")
+
+
+def _capture_typing_extension_shadows(name: str, *collections) -> None:
+    """
+    Ensure different origins for types in typing_extensions are captured.
+    """
+    current_module = sys.modules[__name__]
+    typehint = getattr(current_module, name)
+    if getattr(typehint, "__module__", None) == "typing_extensions" and hasattr(__import__("typing"), name):
+        for collection in collections:
+            collection.add(getattr(__import__("typing"), name))
 
 
 root_types = {
@@ -124,6 +140,8 @@ root_types = {
     OrderedDict,
     Callable,
     abc.Callable,
+    NotRequired,
+    Required,
 }
 
 leaf_types = {
@@ -160,9 +178,20 @@ mapping_origin_types = {
 callable_origin_types = {Callable, abc.Callable}
 
 literal_types = {Literal}
-if getattr(Literal, "__module__", None) == "typing_extensions" and hasattr(__import__("typing"), "Literal"):
-    root_types.add(__import__("typing").Literal)
-    literal_types.add(__import__("typing").Literal)
+_capture_typing_extension_shadows("Literal", root_types, literal_types)
+
+not_required_types = {NotRequired}
+_capture_typing_extension_shadows("NotRequired", root_types, not_required_types)
+
+required_types = {Required}
+_capture_typing_extension_shadows("Required", root_types, required_types)
+not_required_required_types = not_required_types.union(required_types)
+
+typed_dict_types = {TypedDict}
+_capture_typing_extension_shadows("TypedDict", typed_dict_types)
+
+typed_dict_meta_types = {_TypedDictMeta}
+_capture_typing_extension_shadows("_TypedDictMeta", typed_dict_meta_types)
 
 subclass_arg_parser: ContextVar = ContextVar("subclass_arg_parser")
 allow_default_instance: ContextVar = ContextVar("allow_default_instance", default=False)
@@ -889,11 +918,18 @@ def adapt_typehints(
                     else:
                         kwargs["prev_val"] = None
                 val[k] = adapt_typehints(v, subtypehints[1], **kwargs)
-        if get_import_path(typehint.__class__) == "typing._TypedDictMeta":
+        if type(typehint) in typed_dict_meta_types:
             if typehint.__total__:
-                missing_keys = typehint.__annotations__.keys() - val.keys()
-                if missing_keys:
-                    raise_unexpected_value(f"Missing required keys: {missing_keys}", val)
+                required_keys = {
+                    k for k, v in typehint.__annotations__.items() if get_typehint_origin(v) not in not_required_types
+                }
+            else:
+                required_keys = {
+                    k for k, v in typehint.__annotations__.items() if get_typehint_origin(v) in required_types
+                }
+            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()
             if extra_keys:
                 raise_unexpected_value(f"Unexpected keys: {extra_keys}", val)
@@ -904,6 +940,11 @@ def adapt_typehints(
         elif typehint_origin is OrderedDict:
             val = dict(val) if serialize else OrderedDict(val)
 
+    # TypedDict NotRequired and Required
+    elif typehint_origin in not_required_required_types:
+        assert len(subtypehints) == 1, "(Not)Required requires a single type argument"
+        val = adapt_typehints(val, subtypehints[0], **adapt_kwargs)
+
     # Callable
     elif typehint_origin in callable_origin_types or typehint in callable_origin_types:
         if serialize:
diff --git a/jsonargparse/_util.py b/jsonargparse/_util.py
index 44e53d9..e97ea2a 100644
--- a/jsonargparse/_util.py
+++ b/jsonargparse/_util.py
@@ -273,7 +273,7 @@ def get_typehint_origin(typehint):
         typehint_class = get_import_path(typehint.__class__)
         if typehint_class == "types.UnionType":
             return Union
-        if typehint_class == "typing._TypedDictMeta":
+        if typehint_class in {"typing._TypedDictMeta", "typing_extensions._TypedDictMeta"}:
             return dict
     return getattr(typehint, "__origin__", None)
 


Test patch:
diff --git a/jsonargparse_tests/test_typehints.py b/jsonargparse_tests/test_typehints.py
index c391a76..f8f2cde 100644
--- a/jsonargparse_tests/test_typehints.py
+++ b/jsonargparse_tests/test_typehints.py
@@ -24,9 +24,6 @@ from typing import (
     Type,
     Union,
 )
-
-if sys.version_info >= (3, 8):
-    from typing import TypedDict
 from unittest import mock
 from warnings import catch_warnings
 
@@ -37,6 +34,9 @@ from jsonargparse import ArgumentError, Namespace, lazy_instance
 from jsonargparse._typehints import (
     ActionTypeHint,
     Literal,
+    NotRequired,
+    Required,
+    TypedDict,
     get_all_subclass_paths,
     get_subclass_types,
     is_optional,
@@ -501,7 +501,7 @@ def test_mapping_nested_without_args(parser):
     assert {"b": {"c": 2}} == parser.parse_args(['--map={"b": {"c": 2}}']).map
 
 
-@pytest.mark.skipif(sys.version_info < (3, 8), reason="TypedDict introduced in python 3.8")
+@pytest.mark.skipif(not TypedDict, reason="TypedDict introduced in python 3.8 or backported in typing_extensions")
 def test_typeddict_without_arg(parser):
     parser.add_argument("--typeddict", type=TypedDict("MyDict", {}))
     assert {} == parser.parse_args(["--typeddict={}"])["typeddict"]
@@ -513,7 +513,7 @@ def test_typeddict_without_arg(parser):
     ctx.match("Expected a <class 'dict'>")
 
 
-@pytest.mark.skipif(sys.version_info < (3, 8), reason="TypedDict introduced in python 3.8")
+@pytest.mark.skipif(not TypedDict, reason="TypedDict introduced in python 3.8 or backported in typing_extensions")
 def test_typeddict_with_args(parser):
     parser.add_argument("--typeddict", type=TypedDict("MyDict", {"a": int}))
     assert {"a": 1} == parser.parse_args(["--typeddict={'a': 1}"])["typeddict"]
@@ -532,7 +532,7 @@ def test_typeddict_with_args(parser):
     ctx.match("Expected a <class 'dict'>")
 
 
-@pytest.mark.skipif(sys.version_info < (3, 8), reason="TypedDict introduced in python 3.8")
+@pytest.mark.skipif(not TypedDict, reason="TypedDict introduced in python 3.8 or backported in typing_extensions")
 def test_typeddict_with_args_ntotal(parser):
     parser.add_argument("--typeddict", type=TypedDict("MyDict", {"a": int}, total=False))
     assert {"a": 1} == parser.parse_args(["--typeddict={'a': 1}"])["typeddict"]
@@ -548,6 +548,60 @@ def test_typeddict_with_args_ntotal(parser):
     ctx.match("Expected a <class 'dict'>")
 
 
+@pytest.mark.skipif(not NotRequired, reason="NotRequired introduced in python 3.11 or backported in typing_extensions")
+def test_not_required_support(parser):
+    assert ActionTypeHint.is_supported_typehint(NotRequired[Any])
+
+
+@pytest.mark.skipif(not NotRequired, reason="NotRequired introduced in python 3.11 or backported in typing_extensions")
+def test_typeddict_with_not_required_arg(parser):
+    parser.add_argument("--typeddict", type=TypedDict("MyDict", {"a": int, "b": NotRequired[int]}))
+    assert {"a": 1} == parser.parse_args(["--typeddict={'a': 1}"])["typeddict"]
+    assert {"a": 1, "b": 2} == parser.parse_args(["--typeddict={'a': 1, 'b': 2}"])["typeddict"]
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"a":1, "b":2, "c": 3}'])
+    ctx.match("Unexpected keys")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"b":2}'])
+    ctx.match("Missing required keys")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(["--typeddict={}"])
+    ctx.match("Missing required keys")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"a":"x"}'])
+    ctx.match("Expected a <class 'int'>")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"a":1, "b":"x"}'])
+    ctx.match("Expected a <class 'int'>")
+
+
+@pytest.mark.skipif(not Required, reason="Required introduced in python 3.11 or backported in typing_extensions")
+def test_required_support(parser):
+    assert ActionTypeHint.is_supported_typehint(Required[Any])
+
+
+@pytest.mark.skipif(not Required, reason="Required introduced in python 3.11 or backported in typing_extensions")
+def test_typeddict_with_required_arg(parser):
+    parser.add_argument("--typeddict", type=TypedDict("MyDict", {"a": Required[int], "b": int}, total=False))
+    assert {"a": 1} == parser.parse_args(["--typeddict={'a': 1}"])["typeddict"]
+    assert {"a": 1, "b": 2} == parser.parse_args(["--typeddict={'a': 1, 'b': 2}"])["typeddict"]
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"a":1, "b":2, "c": 3}'])
+    ctx.match("Unexpected keys")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"b":2}'])
+    ctx.match("Missing required keys")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(["--typeddict={}"])
+    ctx.match("Missing required keys")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"a":"x"}'])
+    ctx.match("Expected a <class 'int'>")
+    with pytest.raises(ArgumentError) as ctx:
+        parser.parse_args(['--typeddict={"a":1, "b":"x"}'])
+    ctx.match("Expected a <class 'int'>")
+
+
 def test_mapping_proxy_type(parser):
     parser.add_argument("--mapping", type=MappingProxyType)
     cfg = parser.parse_args(['--mapping={"x":1}'])