Try empty context when assigning to union typed variables

mypy/checker.py

@@ -3870,6 +3870,27 @@ def check_simple_assignment(
             rvalue_type = self.expr_checker.accept(
                 rvalue, lvalue_type, always_allow_any=always_allow_any
             )
+            if isinstance(get_proper_type(lvalue_type), UnionType):
+                # Try re-inferring r.h.s. in empty context, and use that if it
+                # results in a narrower type. We don't do this always because this
+                # may cause some perf impact, plus we want to partially preserve
+                # the old behavior. This helps with various practical examples, see
+                # e.g. testOptionalTypeNarrowedByGenericCall.
+                with self.msg.filter_errors() as local_errors, self.local_type_map() as type_map:
+                    alt_rvalue_type = self.expr_checker.accept(
+                        rvalue, None, always_allow_any=always_allow_any
+                    )
+                if (
+                    not local_errors.has_new_errors()
+                    # Skip literal types, as they have special logic (for better errors).
+                    and not isinstance(get_proper_type(rvalue_type), LiteralType)
+                    # Skip Any type, since it is special cased in binder.
+                    and not isinstance(get_proper_type(alt_rvalue_type), AnyType)
+                    and is_valid_inferred_type(alt_rvalue_type)
+                    and is_proper_subtype(alt_rvalue_type, rvalue_type)
+                ):
+                    rvalue_type = alt_rvalue_type
+                    self.store_types(type_map)
             if isinstance(rvalue_type, DeletedType):
                 self.msg.deleted_as_rvalue(rvalue_type, context)
             if isinstance(lvalue_type, DeletedType):

test-data/unit/check-inference-context.test

@@ -1419,3 +1419,60 @@ def bar(x: Union[Mapping[Any, Any], Dict[Any, Sequence[Any]]]) -> None:
     ...
 bar({1: 2})
 [builtins fixtures/dict.pyi]
+
+[case testOptionalTypeNarrowedByGenericCall]
+# flags: --strict-optional
+from typing import Dict, Optional
+
+d: Dict[str, str] = {}
+
+def foo(arg: Optional[str] = None) -> None:
+    if arg is None:
+        arg = d.get("a", "b")
+    reveal_type(arg)  # N: Revealed type is "builtins.str"
+[builtins fixtures/dict.pyi]
+
+[case testOptionalTypeNarrowedByGenericCall2]
+# flags: --strict-optional
+from typing import Dict, Optional
+
+d: Dict[str, str] = {}
+x: Optional[str]
+if x:
+    reveal_type(x)  # N: Revealed type is "builtins.str"
+    x = d.get(x, x)
+    reveal_type(x)  # N: Revealed type is "builtins.str"
+[builtins fixtures/dict.pyi]
+
+[case testOptionalTypeNarrowedByGenericCall3]
+# flags: --strict-optional
+from typing import Generic, TypeVar, Union
+
+T = TypeVar("T")
+def bar(arg: Union[str, T]) -> Union[str, T]: ...
+
+def foo(arg: Union[str, int]) -> None:
+    if isinstance(arg, int):
+        arg = bar("default")
+    reveal_type(arg)  # N: Revealed type is "builtins.str"
+[builtins fixtures/isinstance.pyi]
+
+[case testOptionalTypeNarrowedByGenericCall4]
+# flags: --strict-optional
+from typing import Optional, List, Generic, TypeVar
+
+T = TypeVar("T", covariant=True)
+class C(Generic[T]): ...
+
+x: Optional[C[int]] = None
+y = x = C()
+reveal_type(y)  # N: Revealed type is "__main__.C[builtins.int]"
+
+[case testOptionalTypeNarrowedByGenericCall5]
+from typing import Any, Tuple, Union
+
+i: Union[Tuple[Any, ...], int]
+b: Any
+i = i if isinstance(i, int) else b
+reveal_type(i)  # N: Revealed type is "Union[Any, builtins.int]"
+[builtins fixtures/isinstance.pyi]

test-data/unit/check-typeddict.test

@@ -893,7 +893,7 @@ 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]
+[builtins fixtures/dict.pyi]
 
 [case testTypedDictUnionAmbiguousCase]
 from typing import Union, Mapping, Any, cast