From 46ce325c9ac613985ce96224e45f5c90e5cfd975 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Tue, 22 Jun 2021 13:38:34 +0100
Subject: [PATCH] Fix crash when inferring multiple assignment with overloaded
 function (#10689)

When using lvalue context to re-infer call to an overloaded function,
the inferred tuple type can switch to Any. Defensively accept this. It
probably means that an Any component in argument types causes ambiguity.

Fixes #10653.
---
 mypy/checker.py                       |  9 ++++--
 test-data/unit/check-async-await.test | 41 +++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/mypy/checker.py b/mypy/checker.py
index 556f0ee13104..cd489a2909d6 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -2695,8 +2695,13 @@ def check_multi_assignment_from_tuple(self, lvalues: List[Lvalue], rvalue: Expre
                                                            reinferred_rvalue_type, context,
                                                            infer_lvalue_type)
                     return
-                if isinstance(reinferred_rvalue_type, AnyType) and self.current_node_deferred:
-                    # Doing more inference in deferred nodes can be hard, so give up for now.
+                if isinstance(reinferred_rvalue_type, AnyType):
+                    # We can get Any if the current node is
+                    # deferred. Doing more inference in deferred nodes
+                    # is hard, so give up for now.  We can also get
+                    # here if reinferring types above changes the
+                    # inferred return type for an overloaded function
+                    # to be ambiguous.
                     return
                 assert isinstance(reinferred_rvalue_type, TupleType)
                 rvalue_type = reinferred_rvalue_type
diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test
index e6e5f4be8094..f4bf7c7cfa94 100644
--- a/test-data/unit/check-async-await.test
+++ b/test-data/unit/check-async-await.test
@@ -758,3 +758,44 @@ class Foo(Generic[T]):
 
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
+
+[case testAwaitOverloadSpecialCase]
+from typing import Any, Awaitable, Iterable, overload, Tuple, List, TypeVar, Generic
+
+T = TypeVar("T")
+FT = TypeVar("FT", bound='Future[Any]')
+
+class Future(Awaitable[T], Iterable[T]):
+    pass
+
+class Task(Future[T]):
+    pass
+
+@overload
+def wait(fs: Iterable[FT]) -> Future[Tuple[List[FT], List[FT]]]: ...  \
+    # E: Overloaded function signatures 1 and 2 overlap with incompatible return types
+@overload
+def wait(fs: Iterable[Awaitable[T]]) -> Future[Tuple[List[Task[T]], List[Task[T]]]]: ...
+def wait(fs: Any) -> Any:
+    pass
+
+async def imprecise1(futures: Iterable[Task[Any]]) -> None:
+    done: Any
+    pending: Any
+    done, pending = await wait(futures)
+    reveal_type(done)  # N: Revealed type is "Any"
+
+async def imprecise2(futures: Iterable[Awaitable[Any]]) -> None:
+    done, pending = await wait(futures)
+    reveal_type(done)  # N: Revealed type is "builtins.list[__main__.Task[Any]]"
+
+async def precise1(futures: Iterable[Future[int]]) -> None:
+    done, pending = await wait(futures)
+    reveal_type(done)  # N: Revealed type is "builtins.list[__main__.Future[builtins.int]]"
+
+async def precise2(futures: Iterable[Awaitable[int]]) -> None:
+    done, pending = await wait(futures)
+    reveal_type(done)  # N: Revealed type is "builtins.list[__main__.Task[builtins.int]]"
+
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]