bpo-39627: Fix TypedDict totality check for inherited keys (#18503)

(Adapted from https://github.com/python/typing/pull/700)

3 files changed, 53 insertions, 13 deletions

diff --git a/Lib/test/test_typing.py b/Lib/test/test_typing.py
index bc6a3db..6b0a905 100644
--- a/Lib/test/test_typing.py
+++ b/Lib/test/test_typing.py
@@ -3809,6 +3809,38 @@ class TypedDictTests(BaseTestCase):
         assert Point2Dor3D.__required_keys__ == frozenset(['x', 'y'])
         assert Point2Dor3D.__optional_keys__ == frozenset(['z'])
 
+    def test_keys_inheritance(self):
+        class BaseAnimal(TypedDict):
+            name: str
+
+        class Animal(BaseAnimal, total=False):
+            voice: str
+            tail: bool
+
+        class Cat(Animal):
+            fur_color: str
+
+        assert BaseAnimal.__required_keys__ == frozenset(['name'])
+        assert BaseAnimal.__optional_keys__ == frozenset([])
+        assert BaseAnimal.__annotations__ == {'name': str}
+
+        assert Animal.__required_keys__ == frozenset(['name'])
+        assert Animal.__optional_keys__ == frozenset(['tail', 'voice'])
+        assert Animal.__annotations__ == {
+            'name': str,
+            'tail': bool,
+            'voice': str,
+        }
+
+        assert Cat.__required_keys__ == frozenset(['name', 'fur_color'])
+        assert Cat.__optional_keys__ == frozenset(['tail', 'voice'])
+        assert Cat.__annotations__ == {
+            'fur_color': str,
+            'name': str,
+            'tail': bool,
+            'voice': str,
+        }
+
 
 class IOTests(BaseTestCase):
 
diff --git a/Lib/typing.py b/Lib/typing.py
index 8886b08..6da145f 100644
--- a/Lib/typing.py
+++ b/Lib/typing.py
@@ -1828,23 +1828,30 @@ class _TypedDictMeta(type):
         ns['__new__'] = _typeddict_new if name == 'TypedDict' else _dict_new
         tp_dict = super(_TypedDictMeta, cls).__new__(cls, name, (dict,), ns)
 
-        anns = ns.get('__annotations__', {})
+        annotations = {}
+        own_annotations = ns.get('__annotations__', {})
+        own_annotation_keys = set(own_annotations.keys())
         msg = "TypedDict('Name', {f0: t0, f1: t1, ...}); each t must be a type"
-        anns = {n: _type_check(tp, msg) for n, tp in anns.items()}
-        required = set(anns if total else ())
-        optional = set(() if total else anns)
+        own_annotations = {
+            n: _type_check(tp, msg) for n, tp in own_annotations.items()
+        }
+        required_keys = set()
+        optional_keys = set()
 
         for base in bases:
-            base_anns = base.__dict__.get('__annotations__', {})
-            anns.update(base_anns)
-            if getattr(base, '__total__', True):
-                required.update(base_anns)
-            else:
-                optional.update(base_anns)
+            annotations.update(base.__dict__.get('__annotations__', {}))
+            required_keys.update(base.__dict__.get('__required_keys__', ()))
+            optional_keys.update(base.__dict__.get('__optional_keys__', ()))
+
+        annotations.update(own_annotations)
+        if total:
+            required_keys.update(own_annotation_keys)
+        else:
+            optional_keys.update(own_annotation_keys)
 
-        tp_dict.__annotations__ = anns
-        tp_dict.__required_keys__ = frozenset(required)
-        tp_dict.__optional_keys__ = frozenset(optional)
+        tp_dict.__annotations__ = annotations
+        tp_dict.__required_keys__ = frozenset(required_keys)
+        tp_dict.__optional_keys__ = frozenset(optional_keys)
         if not hasattr(tp_dict, '__total__'):
             tp_dict.__total__ = total
         return tp_dict
diff --git a/Misc/NEWS.d/next/Library/2020-02-13-18-14-15.bpo-39627.Q0scyQ.rst b/Misc/NEWS.d/next/Library/2020-02-13-18-14-15.bpo-39627.Q0scyQ.rst
new file mode 100644
index 0000000..4806aa6
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-02-13-18-14-15.bpo-39627.Q0scyQ.rst
@@ -0,0 +1 @@
+Fixed TypedDict totality check for inherited keys.
\ No newline at end of file