MDS/openide
1
0
Fork 0
mirror of https://gitflic.ru/project/openide/openide.git synced 2026-09-27 10:03:11 +07:00
Code Issues Packages Projects Releases Wiki Activity

PY-41061 Fix type inference for (async) iteration over objects with both __iter__ and __aiter__

Browse Source 
GitOrigin-RevId: 0eb20894a1df40b62b45e7ef0166a7ce72cdc753
This commit is contained in:
Petr
2025-12-02 19:53:12 +00:00
committed by intellij-monorepo-bot
parent 18dcb8c149
commit 39cac66e40
3 changed files with 86 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+33 -25
View File
@@ -218,7 +218,7 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
return PyTypeChecker.getTargetTypeFromTupleAssignment(this, topmostContainingTupleOrList, namedTupleType);
}
else {
return getIterationType(assignedType, assignedIterable, assignedIterable, context);
return getIterationType(assignedType, assignedIterable, assignedIterable, false, context);
}
}
return null;
@@ -312,12 +312,15 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
private @Nullable PyType getTypeFromIteration(@NotNull TypeEvalContext context) {
PyExpression target = null;
PyExpression source = null;
final PyForPart forPart = PsiTreeUtil.getParentOfType(this, PyForPart.class);
if (forPart != null) {
boolean isAsync = false;
PyForStatement forStatement = PsiTreeUtil.getParentOfType(this, PyForStatement.class);
if (forStatement != null) {
final PyForPart forPart = forStatement.getForPart();
final PyExpression expr = forPart.getTarget();
if (PsiTreeUtil.isAncestor(expr, this, false)) {
target = expr;
source = forPart.getSource();
isAsync = forStatement.isAsync();
}
}
final PyComprehensionElement comprh = PsiTreeUtil.getParentOfType(this, PyComprehensionElement.class);
@@ -327,12 +330,13 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
if (PsiTreeUtil.isAncestor(expr, this, false)) {
target = expr;
source = c.getIteratedList();
isAsync = c.isAsync();
}
}
}
if (source != null) {
final PyType sourceType = context.getType(source);
final PyType type = getIterationType(sourceType, source, this, context);
final PyType type = getIterationType(sourceType, source, this, isAsync, context);
target = PyPsiUtils.flattenParens(target);
if (type instanceof PyTupleType tupleType && (target instanceof PyTupleExpression || target instanceof PyListLiteralExpression)) {
return PyTypeChecker.getTargetTypeFromTupleAssignment(this, (PySequenceExpression)target, tupleType);
@@ -346,33 +350,37 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
// TODO migrate this to matching against typing.Iterable protocol with PyTypeUtil.convertToType
public static @Nullable PyType getIterationType(@Nullable PyType iterableType, @Nullable PyExpression source, @NotNull PsiElement anchor,
@NotNull TypeEvalContext context) {
boolean isAsync, @NotNull TypeEvalContext context) {
if (iterableType instanceof PyTupleType tupleType) {
return tupleType.getIteratedItemType();
}
else if (iterableType instanceof PyUnionType) {
return ((PyUnionType)iterableType).map(member -> getIterationType(member, source, anchor, context));
if (iterableType instanceof PyUnionType) {
return ((PyUnionType)iterableType).map(member -> getIterationType(member, source, anchor, isAsync, context));
}
else if (iterableType != null && PyABCUtil.isSubtype(iterableType, PyNames.ITERABLE, context)) {
final PyFunction iterateMethod = findMethodByName(iterableType, PyNames.ITER, context);
if (iterateMethod != null) {
final PyType iterateReturnType = getContextSensitiveType(iterateMethod, context, source);
return getIteratedItemType(iterateReturnType, source, anchor, context, false);
}
final Ref<PyType> nextMethodCallType = getNextMethodCallType(iterableType, source, anchor, context, false);
if (nextMethodCallType != null) {
return nextMethodCallType.get();
}
final PyFunction getItem = findMethodByName(iterableType, PyNames.GETITEM, context);
if (getItem != null) {
return getContextSensitiveType(getItem, context, source);
if (!isAsync) {
if (iterableType != null && PyABCUtil.isSubtype(iterableType, PyNames.ITERABLE, context)) {
final PyFunction iterateMethod = findMethodByName(iterableType, PyNames.ITER, context);
if (iterateMethod != null) {
final PyType iterateReturnType = getContextSensitiveType(iterateMethod, context, source);
return getIteratedItemType(iterateReturnType, source, anchor, context, false);
}
final Ref<PyType> nextMethodCallType = getNextMethodCallType(iterableType, source, anchor, context, false);
if (nextMethodCallType != null) {
return nextMethodCallType.get();
}
final PyFunction getItem = findMethodByName(iterableType, PyNames.GETITEM, context);
if (getItem != null) {
return getContextSensitiveType(getItem, context, source);
}
}
}
else if (iterableType != null && PyABCUtil.isSubtype(iterableType, PyNames.ASYNC_ITERABLE, context)) {
final PyFunction iterateMethod = findMethodByName(iterableType, PyNames.AITER, context);
if (iterateMethod != null) {
final PyType iterateReturnType = getContextSensitiveType(iterateMethod, context, source);
return getIteratedItemType(iterateReturnType, source, anchor, context, true);
else {
if (iterableType != null && PyABCUtil.isSubtype(iterableType, PyNames.ASYNC_ITERABLE, context)) {
final PyFunction iterateMethod = findMethodByName(iterableType, PyNames.AITER, context);
if (iterateMethod != null) {
final PyType iterateReturnType = getContextSensitiveType(iterateMethod, context, source);
return getIteratedItemType(iterateReturnType, source, anchor, context, true);
}
}
}
return null;
python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyYieldExpressionImpl.java
+1 -1
View File
@@ -46,7 +46,7 @@ public class PyYieldExpressionImpl extends PyElementImpl implements PyYieldExpre
final PyType type = expr != null ? context.getType(expr) : PyBuiltinCache.getInstance(this).getNoneType();
if (isDelegating()) {
return PyTargetExpressionImpl.getIterationType(type, expr, this, context);
return PyTargetExpressionImpl.getIterationType(type, expr, this, false, context);
}
return type;
}
python/testSrc/com/jetbrains/python/Py3TypeTest.java
+52
View File
@@ -1071,6 +1071,58 @@ public class Py3TypeTest extends PyTestCase {
""");
}
// PY-41061
public void testForIterationOverObjectWithIterAndAiter() {
doTest("int",
"""
from collections.abc import Iterator, AsyncIterator
class MyIterable:
def __iter__(self) -> Iterator[int]: ...
def __aiter__(self) -> AsyncIterator[str]: ...
for expr in MyIterable(): ...""");
doTest("int",
"""
from collections.abc import Iterator, AsyncIterator
class MyIterable:
def __iter__(self) -> Iterator[int]: ...
def __aiter__(self) -> AsyncIterator[str]: ...
_ = [expr for expr in MyIterable()]""");
}
// PY-41061
public void testAsyncForIterationOverObjectWithIterAndAiter() {
doTest("str",
"""
from collections.abc import Iterator, AsyncIterator
class MyIterable:
def __iter__(self) -> Iterator[int]: ...
def __aiter__(self) -> AsyncIterator[str]: ...
async def foo():
async for expr in MyIterable(): ...""");
doTest("str",
"""
from collections.abc import Iterator, AsyncIterator
class MyIterable:
def __iter__(self) -> Iterator[int]: ...
def __aiter__(self) -> AsyncIterator[str]: ...
async def foo():
_ = [expr async for expr in MyIterable()]""");
}
// PY-21655
public void testUsageOfFunctionDecoratedWithAsyncioCoroutine() {
myFixture.copyDirectoryToProject(TEST_DIRECTORY + getTestName(false), "");