Merge pull request #1201 from google/google_sync

Google sync

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/collections_overlay.py

@@ -34,4 +34,4 @@ class ABCOverlay(typing_overlay.Redirect):
   """A custom overlay for the 'collections.abc' module."""
 
   def __init__(self, ctx):
-    super().__init__("collections.abc", {}, ctx)
+    super().__init__("collections.abc", {"Set": "typing.AbstractSet"}, ctx)

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/overlays/typing_overlay.py

@@ -23,6 +23,10 @@
 Param = overlay_utils.Param
 
 
+def _is_typing_container(cls: pytd.Class):
+  return pytd.IsContainer(cls) and cls.template
+
+
 class TypingOverlay(overlay.Overlay):
   """A representation of the 'typing' module that allows custom overlays.
 
@@ -39,7 +43,7 @@ def __init__(self, ctx):
     ast = ctx.loader.typing
     for cls in ast.classes:
       _, name = cls.name.rsplit(".", 1)
-      if name not in member_map and pytd.IsContainer(cls) and cls.template:
+      if name not in member_map and _is_typing_container(cls):
         member_map[name] = (overlay.build(name, TypingContainer), None)
     super().__init__(ctx, "typing", member_map, ast)
 
@@ -70,7 +74,7 @@ def __init__(self, module_name, aliases, ctx):
     for pyval in ast.aliases + ast.classes + ast.constants + ast.functions:
       # Any public members that are not explicitly implemented are unsupported.
       _, name = pyval.name.rsplit(".", 1)
-      if name.startswith("_"):
+      if name.startswith("_") or name in member_map:
         continue
       if name in typing_overlay:
         member_map[name] = typing_overlay[name][0]
@@ -84,6 +88,9 @@ def __init__(self, module_name, aliases, ctx):
 def _build(name):
   def resolve(ctx):
     ast = ctx.loader.typing
+    pytd_val = ast.Lookup(name)
+    if isinstance(pytd_val, pytd.Class) and _is_typing_container(pytd_val):
+      return TypingContainer(name.rsplit(".", 1)[-1], ctx)
     pytd_type = pytd.ToType(ast.Lookup(name), True, True, True)
     return ctx.convert.constant_to_value(pytd_type)
   return resolve

pytype/tests/test_collections_abc.py

@@ -53,6 +53,21 @@ def f() -> Callable[[], float]: ...
         x: float
       """)
 
+  def test_generator(self):
+    self.Check("""
+      from collections.abc import Generator
+      def f() -> Generator[int, None, None]:
+        yield 0
+    """)
+
+  def test_set(self):
+    # collections.abc.Set is an alias for typing.AbstractSet.
+    self.Check("""
+      from collections.abc import Set
+      def f() -> Set[int]:
+        return frozenset([0])
+    """)
+
 
 if __name__ == "__main__":
   test_base.main()

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)