Raise an exception for typed dict errors rather than logging them immediately.

This lets us filter out bindings that aren't visible.

I also made a couple of tangentially related minor improvements:
* Changed the typed dict error message to say that keys should be "constant
  strings", not just "strings", since the error is also reported when the key
  is a string that can't be resolved to a pyval.
* Silenced a weird lint error.

Resolves #1186.

PiperOrigin-RevId: 447085010

Author: rchen152

pytype/abstract/function.py

@@ -818,8 +818,7 @@ def call_function(ctx,
   if (nodes and not ctx.options.strict_parameter_checks) or not error:
     return node, result
   elif fallback_to_unsolvable:
-    if not isinstance(error, DictKeyMissing):
-      ctx.errorlog.invalid_function_call(ctx.vm.stack(func_var.data[0]), error)
+    ctx.errorlog.invalid_function_call(ctx.vm.stack(func_var.data[0]), error)
     return node, result
   else:
     # We were called by something that does its own error handling.

pytype/errors.py

@@ -963,6 +963,11 @@ def invalid_function_call(self, stack, error):
           stack, error.name, error.bad_call, error.duplicate)
     elif isinstance(error, function.UndefinedParameterError):
       self.name_error(stack, error.name)
+    elif isinstance(error, typed_dict_overlay.TypedDictKeyMissing):
+      self.typed_dict_error(stack, error.typed_dict, error.name)
+    elif isinstance(error, function.DictKeyMissing):
+      # We don't report DictKeyMissing because the false positive rate is high.
+      pass
     else:
       raise AssertionError(error)
 
@@ -1346,7 +1351,8 @@ def typed_dict_error(self, stack, obj, name):
     if name:
       err_msg = f"TypedDict {obj.class_name} does not contain key {name}"
     else:
-      err_msg = f"TypedDict {obj.class_name} requires all keys to be strings"
+      err_msg = (f"TypedDict {obj.class_name} requires all keys to be constant "
+                 "strings")
     self.error(stack, err_msg)
 
   @_error_name("final-error")

pytype/overlays/typed_dict.py

@@ -2,7 +2,7 @@
 
 import dataclasses
 
-from typing import Any, Dict, Set
+from typing import Any, Dict, Optional, Set
 
 from pytype.abstract import abstract
 from pytype.abstract import abstract_utils
@@ -11,6 +11,13 @@
 from pytype.pytd import pytd
 
 
+class TypedDictKeyMissing(function.DictKeyMissing):
+
+  def __init__(self, typed_dict: "TypedDict", key: Optional[str]):
+    super().__init__(key)
+    self.typed_dict = typed_dict
+
+
 @dataclasses.dataclass
 class TypedDictProperties:
   """Collection of typed dict properties passed between various stages."""
@@ -29,7 +36,7 @@ def optional(self):
     return self.keys - self.required
 
   def add(self, k, v, total):
-    self.fields[k] = v
+    self.fields[k] = v  # pylint: disable=unsupported-assignment-operation
     if total:
       self.required.add(k)
 
@@ -220,13 +227,10 @@ def __repr__(self):
 
   def _check_str_key(self, name):
     if name not in self.fields:
-      self.ctx.errorlog.typed_dict_error(self.ctx.vm.frames, self, name)
-      return False
-    return True
+      raise TypedDictKeyMissing(self, name)
 
   def _check_str_key_value(self, node, name, value_var):
-    if not self._check_str_key(name):
-      return
+    self._check_str_key(name)
     typ = abstract_utils.get_atomic_value(self.fields[name])
     bad, _ = self.ctx.matcher(node).bad_matches(value_var, typ)
     for view, error_details in bad:
@@ -240,10 +244,9 @@ def _check_key(self, name_var):
     """Check that key is in the typed dict."""
     try:
       name = abstract_utils.get_atomic_python_constant(name_var, str)
-    except abstract_utils.ConversionError:
-      self.ctx.errorlog.typed_dict_error(self.ctx.vm.frames, self, name=None)
-      return False
-    return self._check_str_key(name)
+    except abstract_utils.ConversionError as e:
+      raise TypedDictKeyMissing(self, None) from e
+    self._check_str_key(name)
 
   def _check_value(self, node, name_var, value_var):
     """Check that value has the right type."""
@@ -254,13 +257,12 @@ def _check_value(self, node, name_var, value_var):
   def getitem_slot(self, node, name_var):
     # A typed dict getitem should have a concrete string arg. If we have a var
     # with multiple bindings just fall back to Any.
-    if not self._check_key(name_var):
-      return node, self.ctx.new_unsolvable(node)
+    self._check_key(name_var)
     return super().getitem_slot(node, name_var)
 
   def setitem_slot(self, node, name_var, value_var):
-    if self._check_key(name_var):
-      self._check_value(node, name_var, value_var)
+    self._check_key(name_var)
+    self._check_value(node, name_var, value_var)
     return super().setitem_slot(node, name_var, value_var)
 
   def set_str_item(self, node, name, value_var):

pytype/tests/test_typed_dict.py

@@ -297,6 +297,23 @@ def f(x: Foo):
       f(x)
     """)
 
+  def test_key_existence_check(self):
+    self.Check("""
+      from typing import Union
+      from typing_extensions import TypedDict
+
+      class Foo(TypedDict):
+        a: int
+      class Bar(TypedDict):
+        b: str
+      class Baz(TypedDict):
+        c: Union[Foo, Bar]
+
+      baz: Baz = {'c': {'a': 0}}
+      assert 'a' in baz['c']
+      print(baz['c']['a'])
+    """)
+
 
 _SINGLE = """
   from typing import TypedDict

pytype/vm_utils.py

@@ -736,8 +736,6 @@ def call_binary_operator(state, name, x, y, report_errors, ctx):
         if ctx.options.report_errors:
           ctx.errorlog.unsupported_operands(ctx.vm.frames, name, x, y)
         result = ctx.new_unsolvable(state.node)
-    elif isinstance(error, function.DictKeyMissing):
-      state, result = error.get_return(state)
     else:
       if ctx.options.report_errors:
         ctx.errorlog.invalid_function_call(ctx.vm.frames, error)