Repository: python/mypy

Issue:
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**



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.

Resolving patch:
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
--- 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]',
@@ -3229,7 +3248,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,
@@ -3295,19 +3314,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
--- a/mypy/messages.py
+++ b/mypy/messages.py
@@ -1128,6 +1128,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:
+        formatted_types = ', '.join(list(format_type_distinctly(*types)))
+        self.fail('Type of TypedDict is ambiguous, could be any of ({})'.format(
+                  formatted_types), context)
+
     def typeddict_key_cannot_be_deleted(
             self,
             typ: TypedDictType,


Test patch:
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
--- a/test-data/unit/check-typeddict.test
+++ b/test-data/unit/check-typeddict.test
@@ -822,6 +822,27 @@ 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})]'
+[builtins fixtures/tuple.pyi]
+
+[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", "B") \
+                                                  # E: Incompatible types in assignment (expression has type "Dict[str, str]", variable has type "Union[A, B]")
+[builtins fixtures/dict.pyi]
 
 -- Use dict literals