Report an error when a TypedDict key has an ambiguous type.

rchen152 · rchen152 · commit bb0736aa4244 · 2024-01-17T17:02:44.000-08:00
* Reports an [invalid-annotation] error when a TypedDict key has an ambiguous type. * Simplifies some of the TypedDict code by assuming that keys have unambiguous types. Fixes #1566. PiperOrigin-RevId: 599320573
diff --git a/pytype/matcher.py b/pytype/matcher.py
@@ -1466,7 +1466,7 @@ def _match_dict_against_typed_dict(
     for k, v in left.pyval.items():
       if k not in fields:
         continue
-      typ = abstract_utils.get_atomic_value(fields[k])
+      typ = fields[k]
       match_result = self.compute_one_match(v, typ)
       if not match_result.success:
         bad.append((k, match_result.bad_matches))
diff --git a/pytype/output.py b/pytype/output.py
@@ -1016,10 +1016,8 @@ def _typed_dict_to_def(self, node, v, name):
       keywords.append(("total", pytd.Literal(False)))
     bases = (pytd.NamedType("typing.TypedDict"),)
     constants = []
-    for k, var in v.props.fields.items():
-      typ = pytd_utils.JoinTypes(
-          self.value_instance_to_pytd_type(node, p, None, set(), {})
-          for p in var.data)
+    for k, val in v.props.fields.items():
+      typ = self.value_instance_to_pytd_type(node, val, None, set(), {})
       if v.props.total and k not in v.props.required:
         typ = pytd.GenericType(pytd.NamedType("typing.NotRequired"), (typ,))
       elif not v.props.total and k in v.props.required:
diff --git a/pytype/overlays/typed_dict.py b/pytype/overlays/typed_dict.py
@@ -2,7 +2,7 @@
 
 import dataclasses
 
-from typing import Any, Dict, Optional, Set
+from typing import Dict, Optional, Set
 
 from pytype.abstract import abstract
 from pytype.abstract import abstract_utils
@@ -34,7 +34,7 @@ class TypedDictProperties:
   """Collection of typed dict properties passed between various stages."""
 
   name: str
-  fields: Dict[str, Any]
+  fields: Dict[str, abstract.BaseValue]
   required: Set[str]
   total: bool
 
@@ -48,27 +48,15 @@ def optional(self):
 
   def add(self, k, v, total):
     """Adds key and value."""
-    values = []
-    all_requiredness = set()
-    for value in v.data:
-      req = _is_required(value)
-      if req is None:
-        values.append(value)
-        all_requiredness.add(None)
-      elif isinstance(value, abstract.ParameterizedClass):
-        values.append(value.formal_type_parameters[abstract_utils.T])
-        all_requiredness.add(req)
-      else:
-        values.append(value.ctx.convert.unsolvable)
-        all_requiredness.add(req)
-    if (len(all_requiredness) == 1 and
-        (requiredness := next(iter(all_requiredness))) is not None):
-      final_v = v.program.NewVariable(values, [], v.program.entrypoint)
-      required = requiredness
+    req = _is_required(v)
+    if req is None:
+      value = v
+    elif isinstance(v, abstract.ParameterizedClass):
+      value = v.formal_type_parameters[abstract_utils.T]
     else:
-      final_v = v
-      required = total
-    self.fields[k] = final_v  # pylint: disable=unsupported-assignment-operation
+      value = v.ctx.convert.unsolvable
+    required = total if req is None else req
+    self.fields[k] = value  # pylint: disable=unsupported-assignment-operation
     if required:
       self.required.add(k)
 
@@ -122,7 +110,12 @@ def _extract_args(self, args):
         name=name, fields={}, required=set(), total=total)
     # Force Required/NotRequired evaluation
     for k, v in fields.items():
-      props.add(k, v, total)
+      try:
+        value = abstract_utils.get_atomic_value(v)
+      except abstract_utils.ConversionError:
+        self.ctx.errorlog.ambiguous_annotation(self.ctx.vm.frames, v.data, k)
+        value = self.ctx.convert.unsolvable
+      props.add(k, value, total)
     return props
 
   def _validate_bases(self, cls_name, bases):
@@ -182,8 +175,14 @@ def make_class(self, node, bases, f_locals, total):
         ordering=classgen.Ordering.FIRST_ANNOTATE,
         ctx=self.ctx)
     for k, local in cls_locals.items():
-      assert local.typ
-      props.add(k, local.typ, total)
+      var = local.typ
+      assert var
+      try:
+        typ = abstract_utils.get_atomic_value(var)
+      except abstract_utils.ConversionError:
+        self.ctx.errorlog.ambiguous_annotation(self.ctx.vm.frames, var.data, k)
+        typ = self.ctx.convert.unsolvable
+      props.add(k, typ, total)
 
     # Process base classes and generate the __init__ signature.
     self._validate_bases(cls_name, bases)
@@ -207,7 +206,7 @@ def make_class_from_pyi(self, cls_name, pytd_cls):
         name=name, fields={}, required=set(), total=total)
 
     for c in pytd_cls.constants:
-      typ = self.ctx.convert.constant_to_var(c.type)
+      typ = self.ctx.convert.constant_to_value(c.type)
       props.add(c.name, typ, total)
 
     # Process base classes and generate the __init__ signature.
@@ -239,8 +238,7 @@ def _make_init(self, props):
     sig = function.Signature.from_param_names(
         f"{props.name}.__init__", props.fields.keys(),
         kind=pytd.ParameterKind.KWONLY)
-    sig.annotations = {k: abstract_utils.get_atomic_value(v)
-                       for k, v in props.fields.items()}
+    sig.annotations = dict(props.fields)
     sig.defaults = {k: self.ctx.new_unsolvable(self.ctx.root_node)
                     for k in props.optional}
     return abstract.SimpleFunction(sig, self.ctx)
@@ -256,8 +254,7 @@ def _new_instance(self, container, node, args):
   def instantiate_value(self, node, container):
     args = function.Args(())
     for name, typ in self.props.fields.items():
-      args.namedargs[name] = self.ctx.join_variables(
-          node, [t.instantiate(node) for t in typ.data])
+      args.namedargs[name] = typ.instantiate(node)
     return self._new_instance(container, node, args)
 
   def instantiate(self, node, container=None):
@@ -301,7 +298,7 @@ def _check_str_key(self, name):
 
   def _check_str_key_value(self, node, name, value_var):
     self._check_str_key(name)
-    typ = abstract_utils.get_atomic_value(self.fields[name])
+    typ = self.fields[name]
     bad = self.ctx.matcher(node).compute_one_match(value_var, typ).bad_matches
     for match in bad:
       self.ctx.errorlog.annotation_type_mismatch(
diff --git a/pytype/tests/test_typed_dict.py b/pytype/tests/test_typed_dict.py
@@ -385,6 +385,14 @@ def f() -> TD:
           return __any_object__
       """)
 
+  def test_duplicate_key(self):
+    self.CheckWithErrors("""
+      from typing_extensions import TypedDict
+      class TD(TypedDict):  # invalid-annotation
+        x: int
+        x: str
+    """)
+
 
 class TypedDictFunctionalTest(test_base.BaseTest):
   """Tests for typing.TypedDict functional constructor."""
@@ -458,6 +466,16 @@ class X(TypedDict, total=False):
         name: str
     """)
 
+  def test_ambiguous_field_type(self):
+    self.CheckWithErrors("""
+      from typing_extensions import TypedDict
+      if __random__:
+        v = str
+      else:
+        v = int
+      X = TypedDict('X', {'k': v})  # invalid-annotation
+    """)
+
 
 _SINGLE = """
   from typing import TypedDict
