Repository: python/mypy (Python)

Issue:
Type narrowing failure inside `AsyncExitStack` blocks
**Bug Report**

In the following synchronous code, mypy is correctly able to perform type narrowing to determine that the `resource` variable must be an instance of `Resource` when `resource.do_thing()` is called inside the `foo()` function:

```python
from contextlib import AbstractContextManager
from typing import TypeVar
T = TypeVar("T")
Self = TypeVar("Self")

class ExitStack:
    def __enter__(self: Self) -> Self:
        return self
    def __exit__(self, *args: object) -> None:
        pass
    def enter_context(self, cm: AbstractContextManager[T]) -> T:
        return cm.__enter__()

class Resource:
    def __enter__(self: Self) -> Self:
        return self
    def __exit__(self, *args: object) -> None:
        pass
    def do_thing(self) -> None:
        pass

def foo(resource: Resource | None) -> None:
    with ExitStack() as stack:
        if resource is None:
            resource = stack.enter_context(Resource())
        resource.do_thing()
```

However, in an asynchronous version of this code, mypy is not able to narrow the type in the same way, leading to a false positive being emitted:

```python
from contextlib import AbstractAsyncContextManager
from typing import TypeVar
T = TypeVar("T")
Self = TypeVar("Self")

class AsyncExitStack:
    async def __aenter__(self: Self) -> Self:
        return self
    async def __aexit__(self, *args: object) -> None:
        pass
    async def enter_async_context(self, cm: AbstractAsyncContextManager[T]) -> T:
        return await cm.__aenter__()

class AsyncResource:
    async def __aenter__(self: Self) -> Self:
        return self
    async def __aexit__(self, *args: object) -> None:
        pass
    def do_thing(self) -> None:
        pass

async def async_foo(resource: AsyncResource | None) -> None:
    async with AsyncExitStack() as stack:
        if resource is None:
            resource = await stack.enter_async_context(AsyncResource())
        resource.do_thing()  # error: Item "None" of "Optional[AsyncResource]" has no attribute "do_thing"
```

Curiously, if we rewrite `async_foo()` like this, the false-positive error goes away:

```python
async def async_foo(resource: AsyncResource | None) -> None:
    async with AsyncExitStack() as stack:
        if resource is None:
            res = await stack.enter_async_context(AsyncResource())
            resource = res
        resource.do_thing()
```

**Real-life use case**

My use case is to write a function like this using `aiohttp`, in which a user can optionally supply an existing `aiohttp.ClientSession()` instance if they want to, or let the function create one on the fly if not.

```python
import aiohttp
from contextlib import AsyncExitStack

async def foo(session: aiohttp.ClientSession | None = None):
    async with AsyncExitStack() as stack:
        if session is None:
            session = await stack.enter_async_context(aiohttp.ClientSession())
        async with session.get("https://foo.com/json") as response:
            data = await response.json()
```

**To Reproduce**

https://mypy-play.net/?mypy=latest&python=3.10&gist=0ee494765b5abc08fc9f161bc63656a7

**Your Environment**

- Mypy version used: 0.982; also occurs on mypy master (commit 8d5b64140c5fe43a2e437e44a6bc3dfdb2b6948a)



Hints:
```async def async_foo(resource: AsyncResource ) -> None:```
can we not remove the or operator here and later resource can be told if it is null

Patch:
diff --git a/mypy/checker.py b/mypy/checker.py
--- a/mypy/checker.py
+++ b/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 patch:
diff --git a/test-data/unit/check-inference-context.test b/test-data/unit/check-inference-context.test
--- a/test-data/unit/check-inference-context.test
+++ b/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]
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
@@ -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