Try empty context when assigning to union typed variables

mypy/checker.py

@@ -76,6 +76,7 @@
     AssignmentStmt,
     Block,
     BreakStmt,
+    BytesExpr,
     CallExpr,
     ClassDef,
     ComparisonExpr,
@@ -86,6 +87,7 @@
     EllipsisExpr,
     Expression,
     ExpressionStmt,
+    FloatExpr,
     ForStmt,
     FuncBase,
     FuncDef,
@@ -115,6 +117,7 @@
     ReturnStmt,
     StarExpr,
     Statement,
+    StrExpr,
     SymbolNode,
     SymbolTable,
     SymbolTableNode,
@@ -3849,6 +3852,23 @@ def inference_error_fallback_type(self, type: Type) -> Type:
         # we therefore need to erase them.
         return erase_typevars(fallback)
 
+    def simple_rvalue(self, rvalue: Expression) -> bool:
+        """Returns True for expressions for which inferred type should not depend on context.
+
+        Note that this function can still return False for some expressions where inferred type
+        does not depend on context. It only exists for performance optimizations.
+        """
+        if isinstance(rvalue, (IntExpr, StrExpr, BytesExpr, FloatExpr, RefExpr)):
+            return True
+        if isinstance(rvalue, CallExpr):
+            if isinstance(rvalue.callee, RefExpr) and isinstance(rvalue.callee.node, FuncBase):
+                typ = rvalue.callee.node.type
+                if isinstance(typ, CallableType):
+                    return not typ.variables
+                elif isinstance(typ, Overloaded):
+                    return not any(item.variables for item in typ.items)
+        return False
+
     def check_simple_assignment(
         self,
         lvalue_type: Type | None,
@@ -3870,6 +3890,30 @@ 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)
+                # Skip literal types, as they have special logic (for better errors).
+                and not isinstance(get_proper_type(rvalue_type), LiteralType)
+                and not self.simple_rvalue(rvalue)
+            ):
+                # 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 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