Repository: jcrist/msgspec (python)

Issue:
Support `NotRequired` in `TypedDict` when `__future__.annotations` is enabled
### Description

[PEP 655](https://peps.python.org/pep-0655/) adds support for `NotRequired` and `Required`. Currently it seems `msgspec` doesn't support it, but it does support `total=False` when defining a `TypedDict`. Is this something that can be supported?

I actually encountered this when I was trying to create a typed dictionary for specifying the values that can be updated in a specific table in Postgres. That is, whichever values are provided would be updated. Initially I had something like this:

```python

class UpdateRequest(TypedDict):

    id: str
    some_value: NotRequired[int]
    another_value: NotRequired[str]
    ...
```
This is what the body of the request from the client was going to be. I refactored this to make it work with msgpec by doing the following:

```python

class UpdateEntity(TypedDict, total=False):

    some_value: int
    another_value: str
    ...

class UpdateRequest(TypedDict):

     id: str
     update_details: UpdateEntity
```

Pull request description:
Support `Required`/`NotRequired` with `__future__.annotations`
Due to a bug in both `typing.TypedDict`/`typing_extensions.TypedDict`, `Required`/`NotRequired` are not properly detected at runtime when `__future__.annotations` is enabled. This means that `TypedDict` definitions making use of these annotations previously couldn't enable `__future__.annotations`. We now parse `Required`/`NotRequired` types manually ourselves, working around the upstream issue.

Fixes #404.

Patch:
diff --git a/msgspec/_utils.py b/msgspec/_utils.py
index a125bf6..c7e1943 100644
--- a/msgspec/_utils.py
+++ b/msgspec/_utils.py
@@ -171,20 +171,33 @@ def get_typeddict_info(obj):
         cls = obj
     else:
         cls = obj.__origin__
-    hints = get_class_annotations(obj)
-    # Strip off Required/NotRequired
-    hints = {
-        k: v.__args__[0]
-        if getattr(v, "__origin__", False) in (Required, NotRequired)
-        else v
-        for k, v in hints.items()
-    }
+
+    raw_hints = get_class_annotations(obj)
+
     if hasattr(cls, "__required_keys__"):
         required = set(cls.__required_keys__)
     elif cls.__total__:
-        required = set(hints)
+        required = set(raw_hints)
     else:
         required = set()
+
+    # Both `typing.TypedDict` and `typing_extensions.TypedDict` have a bug
+    # where `Required`/`NotRequired` aren't properly detected at runtime when
+    # `__future__.annotations` is enabled, meaning the `__required_keys__`
+    # isn't correct. This code block works around this issue by amending the
+    # set of required keys as needed, while also stripping off any
+    # `Required`/`NotRequired` wrappers.
+    hints = {}
+    for k, v in raw_hints.items():
+        origin = getattr(v, "__origin__", False)
+        if origin is Required:
+            required.add(k)
+            hints[k] = v.__args__[0]
+        elif origin is NotRequired:
+            required.discard(k)
+            hints[k] = v.__args__[0]
+        else:
+            hints[k] = v
     return hints, required
 
 


Test patch:
diff --git a/tests/test_common.py b/tests/test_common.py
index 6606631..04f0bf7 100644
--- a/tests/test_common.py
+++ b/tests/test_common.py
@@ -1762,6 +1762,52 @@ class TestTypedDict:
             dec.decode(proto.encode({"b": "two"}))
         assert "Object missing required field `a`" == str(rec.value)
 
+    @pytest.mark.parametrize("use_typing_extensions", [False, True])
+    def test_required_and_notrequired(self, proto, use_typing_extensions):
+        if use_typing_extensions:
+            module = "typing_extensions"
+        else:
+            module = "typing"
+
+        ns = pytest.importorskip(module)
+
+        if not hasattr(ns, "Required"):
+            pytest.skip(f"{module}.Required is not available")
+
+        if not hasattr(ns.TypedDict("C"), "__required_keys__"):
+            # This should be Python 3.8, builtin typing only
+            pytest.skip("partially optional TypedDict not supported")
+
+        source = f"""
+        from __future__ import annotations
+        from {module} import TypedDict, Required, NotRequired
+
+        class Base(TypedDict):
+            a: int
+            b: NotRequired[str]
+
+        class Ex(Base, total=False):
+            c: str
+            d: Required[bool]
+        """
+
+        with temp_module(source) as mod:
+            dec = proto.Decoder(mod.Ex)
+
+            x = {"a": 1, "b": "two", "c": "extra", "d": False}
+            assert dec.decode(proto.encode(x)) == x
+
+            x2 = {"a": 1, "d": False}
+            assert dec.decode(proto.encode(x2)) == x2
+
+            with pytest.raises(msgspec.ValidationError) as rec:
+                dec.decode(proto.encode({"d": False}))
+            assert "Object missing required field `a`" == str(rec.value)
+
+            with pytest.raises(msgspec.ValidationError) as rec:
+                dec.decode(proto.encode({"a": 2}))
+            assert "Object missing required field `d`" == str(rec.value)
+
     def test_keys_are_their_interned_values(self, proto):
         """Ensure that we're not allocating new keys here, but reusing the
         existing keys on the TypedDict schema"""