Repository: quora/pyanalyze
Pull request #567 (2022-11-14 20:07:24)

Error for invalid TypedDict keys


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 4834160..1906a5f 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Emit an error for unknown `TypedDict` keys (#567)
 - Improve type inference for f-strings containing literals (#571)
 - Add experimental `@has_extra_keys` decorator for `TypedDict` types (#568)
 - Fix crash on recursive type aliases. Recursive type aliases now fall back to `Any` (#565)
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index 5297b89..a6f8bde 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -557,14 +557,21 @@ def _dict_getitem_impl(ctx: CallContext) -> ImplReturn:
                 # probably KeyError, but catch anything in case it's an
                 # unhashable str subclass or something
                 except Exception:
-                    # No error here; TypedDicts may have additional keys at runtime.
-                    pass
+                    if self_value.extra_keys is None:
+                        ctx.show_error(
+                            f"Unknown TypedDict key {key}",
+                            ErrorCode.invalid_typeddict_key,
+                            arg="k",
+                        )
+                        return AnyValue(AnySource.error)
             if self_value.extra_keys is not None:
                 return self_value.extra_keys
-            # TODO strictly we should throw an error for any non-Literal or unknown key:
-            # https://www.python.org/dev/peps/pep-0589/#supported-and-unsupported-operations
-            # Don't do that yet because it may cause too much disruption.
-            return AnyValue(AnySource.inference)
+            ctx.show_error(
+                f"TypedDict key must be a literal, not {key}",
+                ErrorCode.invalid_typeddict_key,
+                arg="k",
+            )
+            return AnyValue(AnySource.error)
         elif isinstance(self_value, DictIncompleteValue):
             val = self_value.get_value(key, ctx.visitor)
             if val is UNINITIALIZED_VALUE:
@@ -622,8 +629,13 @@ def _dict_get_impl(ctx: CallContext) -> ImplReturn:
                 # probably KeyError, but catch anything in case it's an
                 # unhashable str subclass or something
                 except Exception:
-                    # No error here; TypedDicts may have additional keys at runtime.
-                    pass
+                    if self_value.extra_keys is None:
+                        ctx.show_error(
+                            f"Unknown TypedDict key {key.val!r}",
+                            ErrorCode.invalid_typeddict_key,
+                            arg="k",
+                        )
+                        return AnyValue(AnySource.error)
                 else:
                     if required:
                         return value
@@ -631,10 +643,12 @@ def _dict_get_impl(ctx: CallContext) -> ImplReturn:
                         return value | default
             if self_value.extra_keys is not None:
                 return self_value.extra_keys | default
-            # TODO strictly we should throw an error for any non-Literal or unknown key:
-            # https://www.python.org/dev/peps/pep-0589/#supported-and-unsupported-operations
-            # Don't do that yet because it may cause too much disruption.
-            return default
+            ctx.show_error(
+                f"TypedDict key must be a literal, not {key}",
+                ErrorCode.invalid_typeddict_key,
+                arg="k",
+            )
+            return AnyValue(AnySource.error)
         elif isinstance(self_value, DictIncompleteValue):
             val = self_value.get_value(key, ctx.visitor)
             if val is UNINITIALIZED_VALUE:


Test patch:
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index 1a6223a..519c877 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -1748,75 +1748,6 @@ class TestFStrings(TestNameCheckVisitorBase):
             return f"{x}"  # E: undefined_name
 
 
-class TestTypedDict(TestNameCheckVisitorBase):
-    @assert_passes()
-    def test_basic(self):
-        from mypy_extensions import TypedDict as METypedDict
-        from typing_extensions import TypedDict as TETypedDict
-
-        T = METypedDict("T", {"a": int, "b": str})
-        T2 = TETypedDict("T2", {"a": int, "b": str})
-
-        def capybara(x: T, y: T2):
-            assert_is_value(
-                x,
-                TypedDictValue(
-                    {"a": (True, TypedValue(int)), "b": (True, TypedValue(str))}
-                ),
-            )
-            assert_is_value(x["a"], TypedValue(int))
-            assert_is_value(
-                y,
-                TypedDictValue(
-                    {"a": (True, TypedValue(int)), "b": (True, TypedValue(str))}
-                ),
-            )
-            assert_is_value(y["a"], TypedValue(int))
-
-    @assert_passes()
-    def test_unknown_key_unresolved(self):
-        from mypy_extensions import TypedDict
-
-        T = TypedDict("T", {"a": int, "b": str})
-
-        def capybara(x: T):
-            assert_is_value(x["not a key"], AnyValue(AnySource.inference))
-
-    @assert_passes()
-    def test_invalid_key(self):
-        from mypy_extensions import TypedDict
-
-        T = TypedDict("T", {"a": int, "b": str})
-
-        def capybara(x: T):
-            x[0]  # E: invalid_typeddict_key
-
-    @assert_passes()
-    def test_total(self):
-        from typing_extensions import TypedDict
-
-        class TD(TypedDict, total=False):
-            a: int
-            b: str
-
-        class TD2(TD):
-            c: float
-
-        def f(td: TD) -> None:
-            pass
-
-        def g(td2: TD2) -> None:
-            pass
-
-        def caller() -> None:
-            f({})
-            f({"a": 1})
-            f({"a": 1, "b": "c"})
-            f({"a": "a"})  # E: incompatible_argument
-            g({"c": 1.0})
-            g({})  # E: incompatible_argument
-
-
 _AnnotSettings = {
     ErrorCode.missing_parameter_annotation: True,
     ErrorCode.missing_return_annotation: True,
diff --git a/pyanalyze/test_typeddict.py b/pyanalyze/test_typeddict.py
index 68279bf..547d046 100644
--- a/pyanalyze/test_typeddict.py
+++ b/pyanalyze/test_typeddict.py
@@ -2,7 +2,7 @@
 from .implementation import assert_is_value
 from .test_name_check_visitor import TestNameCheckVisitorBase
 from .test_node_visitor import assert_passes
-from .value import TypedDictValue, TypedValue
+from .value import TypedDictValue, TypedValue, AnyValue, AnySource
 
 
 class TestExtraKeys(TestNameCheckVisitorBase):
@@ -155,3 +155,82 @@ class TestTypedDict(TestNameCheckVisitorBase):
                     {"x": (True, TypedValue(int)), "y": (False, TypedValue(str))}
                 ),
             )
+
+    @assert_passes()
+    def test_unknown_key(self):
+        from typing_extensions import TypedDict, assert_type
+
+        class Capybara(TypedDict):
+            x: int
+
+        def user(c: Capybara):
+            assert_type(c["x"], int)
+            c["y"]  # E: invalid_typeddict_key
+
+    @assert_passes()
+    def test_basic(self):
+        from mypy_extensions import TypedDict as METypedDict
+        from typing_extensions import TypedDict as TETypedDict
+
+        T = METypedDict("T", {"a": int, "b": str})
+        T2 = TETypedDict("T2", {"a": int, "b": str})
+
+        def capybara(x: T, y: T2):
+            assert_is_value(
+                x,
+                TypedDictValue(
+                    {"a": (True, TypedValue(int)), "b": (True, TypedValue(str))}
+                ),
+            )
+            assert_is_value(x["a"], TypedValue(int))
+            assert_is_value(
+                y,
+                TypedDictValue(
+                    {"a": (True, TypedValue(int)), "b": (True, TypedValue(str))}
+                ),
+            )
+            assert_is_value(y["a"], TypedValue(int))
+
+    @assert_passes()
+    def test_unknown_key_unresolved(self):
+        from mypy_extensions import TypedDict
+
+        T = TypedDict("T", {"a": int, "b": str})
+
+        def capybara(x: T):
+            val = x["not a key"]  # E: invalid_typeddict_key
+            assert_is_value(val, AnyValue(AnySource.error))
+
+    @assert_passes()
+    def test_invalid_key(self):
+        from mypy_extensions import TypedDict
+
+        T = TypedDict("T", {"a": int, "b": str})
+
+        def capybara(x: T):
+            x[0]  # E: invalid_typeddict_key
+
+    @assert_passes()
+    def test_total(self):
+        from typing_extensions import TypedDict
+
+        class TD(TypedDict, total=False):
+            a: int
+            b: str
+
+        class TD2(TD):
+            c: float
+
+        def f(td: TD) -> None:
+            pass
+
+        def g(td2: TD2) -> None:
+            pass
+
+        def caller() -> None:
+            f({})
+            f({"a": 1})
+            f({"a": 1, "b": "c"})
+            f({"a": "a"})  # E: incompatible_argument
+            g({"c": 1.0})
+            g({})  # E: incompatible_argument