Repository: python/mypy
Pull request: Typeddict context returned correctly when unambiguous
Fixes #8156 

This is a suggestion as to how we could handle getting the Typeddict context from a Union with more than one Typeddict. 

The idea is that we check each of the Typeddict types to find one who's keys match that of the expression, if more than one match it's then ambiguous. The slightly weird thing is that this matching is only applied in the Union case, if there is only a single Typeddict it's returned as before. 

I'm not sure that I've put the tests in the most appropriate place, and suggestions would be great. 

Problem statement:
Dictionaries not accepted for implicit cast to TypedDict Union
Please provide more information to help us understand the issue:

* Are you reporting a bug, or opening a feature request? **Bug**

```python
from typing import TypedDict, Literal, Union, Mapping, Any, cast

A = TypedDict('A', {'@type': Literal['a-type'], 'a': str})
B = TypedDict('B', {'@type': Literal['b-type'], 'b': int})

bad_cast: Union[A, B] = {'@type': 'a-type', 'a': 'Test'}
good_cast_1: Union[A, B] = cast(Union[A, B], {'@type': 'a-type', 'a': 'Test'})
good_cast_2: Union[A, B] = cast(A, {'@type': 'a-type', 'a': 'Test'})
```

* What is the actual behavior/output?

```test_mypy_discriminate.py:6:25: error: Incompatible types in assignment (expression has type "Dict[str, str]", variable has type "Union[A, B]")```

* What is the behavior/output you expect?

Since the given type matches one of the types in the Union, it should be accepted.

* What are the versions of mypy and Python you are using? **Python 3.8, mypy 0.750**
* Do you see the same issue after installing mypy from Git master? **Don't know how to do that**

Hints from issue discussion:
Yeah, I think this is a bug. The TypedDict return context is special-cased, but as it often happens special-casing doesn't take into account union types.

Commit message: Handle validated_kwargs better in match function

Patch to review:
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 7edaf7e2ad89..ac808fe4ce04 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -487,9 +487,8 @@ def check_typeddict_call(self, callee: TypedDictType,
         self.chk.fail(message_registry.INVALID_TYPEDDICT_ARGS, context)
         return AnyType(TypeOfAny.from_error)
 
-    def check_typeddict_call_with_dict(self, callee: TypedDictType,
-                                       kwargs: DictExpr,
-                                       context: Context) -> Type:
+    def validate_typeddict_kwargs(
+            self, kwargs: DictExpr) -> 'Optional[OrderedDict[str, Expression]]':
         item_args = [item[1] for item in kwargs.items]
 
         item_names = []  # List[str]
@@ -504,12 +503,32 @@ def check_typeddict_call_with_dict(self, callee: TypedDictType,
                 key_context = item_name_expr or item_arg
                 self.chk.fail(message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL,
                               key_context)
-                return AnyType(TypeOfAny.from_error)
+                return None
             else:
                 item_names.append(literal_value)
+        return OrderedDict(zip(item_names, item_args))
+
+    def match_typeddict_call_with_dict(self, callee: TypedDictType,
+                                       kwargs: DictExpr,
+                                       context: Context) -> bool:
+        validated_kwargs = self.validate_typeddict_kwargs(kwargs=kwargs)
+        if validated_kwargs is not None:
+            return (callee.required_keys <= set(validated_kwargs.keys())
+                <= set(callee.items.keys()))
+        else:
+            return False
 
-        return self.check_typeddict_call_with_kwargs(
-            callee, OrderedDict(zip(item_names, item_args)), context)
+    def check_typeddict_call_with_dict(self, callee: TypedDictType,
+                                       kwargs: DictExpr,
+                                       context: Context) -> Type:
+        validated_kwargs = self.validate_typeddict_kwargs(kwargs=kwargs)
+        if validated_kwargs is not None:
+            return self.check_typeddict_call_with_kwargs(
+                callee,
+                kwargs=validated_kwargs,
+                context=context)
+        else:
+            return AnyType(TypeOfAny.from_error)
 
     def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
                                          kwargs: 'OrderedDict[str, Expression]',
@@ -3228,7 +3247,7 @@ def visit_dict_expr(self, e: DictExpr) -> Type:
         # an error, but returns the TypedDict type that matches the literal it found
         # that would cause a second error when that TypedDict type is returned upstream
         # to avoid the second error, we always return TypedDict type that was requested
-        typeddict_context = self.find_typeddict_context(self.type_context[-1])
+        typeddict_context = self.find_typeddict_context(self.type_context[-1], e)
         if typeddict_context:
             self.check_typeddict_call_with_dict(
                 callee=typeddict_context,
@@ -3294,19 +3313,25 @@ def visit_dict_expr(self, e: DictExpr) -> Type:
         assert rv is not None
         return rv
 
-    def find_typeddict_context(self, context: Optional[Type]) -> Optional[TypedDictType]:
+    def find_typeddict_context(self, context: Optional[Type],
+                              dict_expr: DictExpr) -> Optional[TypedDictType]:
         context = get_proper_type(context)
         if isinstance(context, TypedDictType):
             return context
         elif isinstance(context, UnionType):
             items = []
             for item in context.items:
-                item_context = self.find_typeddict_context(item)
-                if item_context:
+                item_context = self.find_typeddict_context(item, dict_expr)
+                if (item_context is not None
+                        and self.match_typeddict_call_with_dict(
+                            item_context, dict_expr, dict_expr)):
                     items.append(item_context)
             if len(items) == 1:
-                # Only one union item is TypedDict, so use the context as it's unambiguous.
+                # Only one union item is valid TypedDict for the given dict_expr, so use the
+                # context as it's unambiguous.
                 return items[0]
+            if len(items) > 1:
+                self.msg.typeddict_context_ambiguous(items, dict_expr)
         # No TypedDict type in context.
         return None
 
diff --git a/mypy/messages.py b/mypy/messages.py
index 6eddc8184cb1..f41a70745df9 100644
--- a/mypy/messages.py
+++ b/mypy/messages.py
@@ -1111,6 +1111,14 @@ def typeddict_key_not_found(
         else:
             self.fail("TypedDict {} has no key '{}'".format(format_type(typ), item_name), context)
 
+    def typeddict_context_ambiguous(
+            self,
+            types: List[TypedDictType],
+            context: Context) -> None:
+        self.fail('Type of TypedDict is Ambiguous, could be any of {}'.format(
+                  ([(format_type(typ), format_item_name_list(typ.items.keys())) for typ in types])
+                  ), context)
+
     def typeddict_key_cannot_be_deleted(
             self,
             typ: TypedDictType,
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
index e87d20cf61a9..3726a1db8424 100644
--- a/test-data/unit/check-typeddict.test
+++ b/test-data/unit/check-typeddict.test
@@ -822,6 +822,26 @@ reveal_type(u(c, m_i_i)) # N: Revealed type is 'Union[typing.Mapping*[builtins.i
 reveal_type(u(c, m_s_a)) # N: Revealed type is 'Union[typing.Mapping*[builtins.str, Any], TypedDict('__main__.C', {'a': builtins.int, 'b': builtins.int})]'
 [builtins fixtures/dict.pyi]
 
+[case testTypedDictUnionUnambiguousCase]
+from typing import Union, Mapping, Any, cast
+from typing_extensions import TypedDict, Literal
+
+A = TypedDict('A', {'@type': Literal['a-type'], 'a': str})
+B = TypedDict('B', {'@type': Literal['b-type'], 'b': int})
+
+c: Union[A, B] = {'@type': 'a-type', 'a': 'Test'}
+reveal_type(c) # N: Revealed type is 'Union[TypedDict('__main__.A', {'@type': Literal['a-type'], 'a': builtins.str}), TypedDict('__main__.B', {'@type': Literal['b-type'], 'b': builtins.int})]'
+
+[case testTypedDictUnionAmbiguousCase]
+from typing import Union, Mapping, Any, cast
+from typing_extensions import TypedDict, Literal
+
+A = TypedDict('A', {'@type': Literal['a-type'], 'a': str})
+B = TypedDict('B', {'@type': Literal['a-type'], 'a': str})
+
+c: Union[A, B] = {'@type': 'a-type', 'a': 'Test'} # E: Type of TypedDict is Ambiguous, could be any of [('"A"', "('@type', 'a')"), ('"B"', "('@type', 'a')")] \
+                                                  # E: Incompatible types in assignment (expression has type "Dict[str, str]", variable has type "Union[A, B]")
+[builtins fixtures/dict.pyi]
 
 -- Use dict literals
 


Code review comments:
diff_hunk: @@ -1111,6 +1111,14 @@ def typeddict_key_not_found(
         else:
             self.fail("TypedDict {} has no key '{}'".format(format_type(typ), item_name), context)
 
+    def typeddict_context_ambiguous(
+            self,
+            types: List[TypedDictType],
+            context: Context) -> None:
+        self.fail('Type of TypedDict is Ambiguous, could be any of {}'.format(
path: mypy/messages.py
text: @user1:
Don't capitalize Ambiguous

diff_hunk: @@ -1111,6 +1111,14 @@ def typeddict_key_not_found(
         else:
             self.fail("TypedDict {} has no key '{}'".format(format_type(typ), item_name), context)
 
+    def typeddict_context_ambiguous(
+            self,
+            types: List[TypedDictType],
+            context: Context) -> None:
+        self.fail('Type of TypedDict is Ambiguous, could be any of {}'.format(
+                  ([(format_type(typ), format_item_name_list(typ.items.keys())) for typ in types])
path: mypy/messages.py
text: @user1:
The type error message here is ugly. My inclination is to just print the types, though if you can make the rest look good and having clear meaning then that's OK. We definitely don't want to be printing out a bunch of brackets and parens and quotes that come from python's default reprs.

If multiple names are being formatted, also, we want to use `format_type_distinctly` to make sure they don't end up getting given the same message. Unfortunately `format_type_distinctly` only works on two strings, so I went and made a PR (#8270) to fix that.

@author:
Yeah I wasn't quite sure what to do about printing the types here. Thanks a lot for that PR, I've updated to use ```format_type_distinctly``` and it's much cleaner! Also if you can think of any slightly nicer wording do let me know.