gh-112509: Fix keys being present in both required_keys and optional_keys in TypedDict (#112512)

JelleZijlstra · AlexWaygood · web-flow · commit 403886942376 · 2023-11-29T09:36:48.000-08:00
Co-authored-by: Alex Waygood &lt;Alex.Waygood@Gmail.com&gt;
diff --git a/Lib/test/test_typing.py b/Lib/test/test_typing.py
@@ -7769,6 +7769,46 @@ class Cat(Animal):
             'voice': str,
         })
 
+    def test_keys_inheritance_with_same_name(self):
+        class NotTotal(TypedDict, total=False):
+            a: int
+
+        class Total(NotTotal):
+            a: int
+
+        self.assertEqual(NotTotal.__required_keys__, frozenset())
+        self.assertEqual(NotTotal.__optional_keys__, frozenset(['a']))
+        self.assertEqual(Total.__required_keys__, frozenset(['a']))
+        self.assertEqual(Total.__optional_keys__, frozenset())
+
+        class Base(TypedDict):
+            a: NotRequired[int]
+            b: Required[int]
+
+        class Child(Base):
+            a: Required[int]
+            b: NotRequired[int]
+
+        self.assertEqual(Base.__required_keys__, frozenset(['b']))
+        self.assertEqual(Base.__optional_keys__, frozenset(['a']))
+        self.assertEqual(Child.__required_keys__, frozenset(['a']))
+        self.assertEqual(Child.__optional_keys__, frozenset(['b']))
+
+    def test_multiple_inheritance_with_same_key(self):
+        class Base1(TypedDict):
+            a: NotRequired[int]
+
+        class Base2(TypedDict):
+            a: Required[str]
+
+        class Child(Base1, Base2):
+            pass
+
+        # Last base wins
+        self.assertEqual(Child.__annotations__, {'a': Required[str]})
+        self.assertEqual(Child.__required_keys__, frozenset(['a']))
+        self.assertEqual(Child.__optional_keys__, frozenset())
+
     def test_required_notrequired_keys(self):
         self.assertEqual(NontotalMovie.__required_keys__,
                          frozenset({"title"}))
diff --git a/Lib/typing.py b/Lib/typing.py
@@ -2884,8 +2884,14 @@ def __new__(cls, name, bases, ns, total=True):
 
         for base in bases:
             annotations.update(base.__dict__.get('__annotations__', {}))
-            required_keys.update(base.__dict__.get('__required_keys__', ()))
-            optional_keys.update(base.__dict__.get('__optional_keys__', ()))
+
+            base_required = base.__dict__.get('__required_keys__', set())
+            required_keys |= base_required
+            optional_keys -= base_required
+
+            base_optional = base.__dict__.get('__optional_keys__', set())
+            required_keys -= base_optional
+            optional_keys |= base_optional
 
         annotations.update(own_annotations)
         for annotation_key, annotation_type in own_annotations.items():
@@ -2897,14 +2903,23 @@ def __new__(cls, name, bases, ns, total=True):
                     annotation_origin = get_origin(annotation_type)
 
             if annotation_origin is Required:
-                required_keys.add(annotation_key)
+                is_required = True
             elif annotation_origin is NotRequired:
-                optional_keys.add(annotation_key)
-            elif total:
+                is_required = False
+            else:
+                is_required = total
+
+            if is_required:
                 required_keys.add(annotation_key)
+                optional_keys.discard(annotation_key)
             else:
                 optional_keys.add(annotation_key)
+                required_keys.discard(annotation_key)
 
+        assert required_keys.isdisjoint(optional_keys), (
+            f"Required keys overlap with optional keys in {name}:"
+            f" {required_keys=}, {optional_keys=}"
+        )
         tp_dict.__annotations__ = annotations
         tp_dict.__required_keys__ = frozenset(required_keys)
         tp_dict.__optional_keys__ = frozenset(optional_keys)
diff --git a/Misc/NEWS.d/next/Library/2023-11-28-20-01-33.gh-issue-112509.QtoKed.rst b/Misc/NEWS.d/next/Library/2023-11-28-20-01-33.gh-issue-112509.QtoKed.rst
@@ -0,0 +1,3 @@
+Fix edge cases that could cause a key to be present in both the
+``__required_keys__`` and ``__optional_keys__`` attributes of a
+:class:`typing.TypedDict`. Patch by Jelle Zijlstra.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+Fix edge cases that could cause a key to be present in both the`
	`2`	+``__required_keys__`` and ``__optional_keys__`` attributes of a
	`3`	+:class:`typing.TypedDict`. Patch by Jelle Zijlstra.