diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java index 4ad878c3a6cd..a7758cae9cc6 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java @@ -381,10 +381,8 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl nextMethodCallType = getNextMethodCallType(type, source, anchor, context, async); if (nextMethodCallType != null && !nextMethodCallType.isNull()) { - if (nextMethodCallType.get() instanceof PyCollectionType collectionType) { - if (async && "typing.Awaitable".equals(collectionType.getClassQName())) { - return collectionType.getIteratedItemType(); - } + if (async) { + return Ref.deref(PyTypingTypeProvider.unwrapCoroutineReturnType(nextMethodCallType.get())); } return nextMethodCallType.get(); } diff --git a/python/testSrc/com/jetbrains/python/Py3TypeTest.java b/python/testSrc/com/jetbrains/python/Py3TypeTest.java index 98346e32bc3a..e259b57e2e16 100644 --- a/python/testSrc/com/jetbrains/python/Py3TypeTest.java +++ b/python/testSrc/com/jetbrains/python/Py3TypeTest.java @@ -568,7 +568,7 @@ public class Py3TypeTest extends PyTestCase { doTest("int", """ class AIter(object): - def __anext__(self): + async def __anext__(self): return 5 class A(object): def __aiter__(self): @@ -577,6 +577,21 @@ public class Py3TypeTest extends PyTestCase { async for expr in a: print(expr)"""); } + + // PY-60714 + public void testAsyncIteratorUnwrapsCoroutineFromAnext() { + doTest("bytes", """ + class AIterator: + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + return b"a" + + async for expr in AIterator(): + print(expt) + """); + } // PY-21655 public void testUsageOfFunctionDecoratedWithAsyncioCoroutine() {