diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyForPartImpl.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyForPartImpl.java index 998900ff526c..120e4d58b9cd 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyForPartImpl.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyForPartImpl.java @@ -2,10 +2,26 @@ package com.jetbrains.python.psi.impl; import com.intellij.lang.ASTNode; +import com.intellij.openapi.util.TextRange; +import com.intellij.psi.PsiReference; import com.jetbrains.python.psi.PyForPart; +import com.jetbrains.python.PyTokenTypes; +import com.jetbrains.python.psi.impl.references.PyKeywordReference; +import com.jetbrains.python.psi.resolve.PyResolveContext; +import com.jetbrains.python.psi.types.TypeEvalContext; public class PyForPartImpl extends PyElementImpl implements PyForPart { public PyForPartImpl(ASTNode astNode) { super(astNode); } + + @Override + public PsiReference getReference() { + var inNode = getNode().findChildByType(PyTokenTypes.IN_KEYWORD); + if (inNode == null) return null; + + TextRange range = inNode.getPsi().getTextRangeInParent(); + PyResolveContext resolveContext = PyResolveContext.defaultContext(TypeEvalContext.codeAnalysis(getProject(), getContainingFile())); + return new PyKeywordReference(this, resolveContext, range); + } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyKeywordReference.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyKeywordReference.java new file mode 100644 index 000000000000..f21fefb4e777 --- /dev/null +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyKeywordReference.java @@ -0,0 +1,67 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. +package com.jetbrains.python.psi.impl.references; + +import com.intellij.openapi.util.TextRange; +import com.intellij.psi.PsiReferenceBase; +import com.intellij.psi.ResolveResult; +import com.jetbrains.python.PyNames; +import com.jetbrains.python.psi.*; +import com.intellij.psi.util.PsiTreeUtil; +import com.jetbrains.python.psi.resolve.PyResolveContext; +import com.jetbrains.python.psi.resolve.RatedResolveResult; +import com.jetbrains.python.psi.AccessDirection; +import com.jetbrains.python.psi.types.PyType; +import com.jetbrains.python.psi.types.TypeEvalContext; +import org.jetbrains.annotations.NotNull; + +import java.util.ArrayList; +import java.util.List; + +public final class PyKeywordReference extends PsiReferenceBase.Poly { + + private final PyResolveContext myContext; + + public PyKeywordReference(@NotNull PyElement owner, @NotNull PyResolveContext context, @NotNull TextRange rangeInOwner) { + super(owner, rangeInOwner, false); + myContext = context; + } + + @Override + public boolean isSoft() { + return true; + } + + @Override + public ResolveResult @NotNull [] multiResolve(boolean incompleteCode) { + final PyElement element = getElement(); + if (element instanceof PyForPart forPart) { + final PyExpression source = forPart.getSource(); + if (source == null) return ResolveResult.EMPTY_ARRAY; + + final TypeEvalContext typeEvalContext = myContext.getTypeEvalContext(); + final PyType type = typeEvalContext.getType(source); + if (type == null) return ResolveResult.EMPTY_ARRAY; + + final List results = new ArrayList<>(); + + boolean isAsync = false; + PyForStatement forStatement = PsiTreeUtil.getParentOfType(forPart, PyForStatement.class); + if (forStatement != null) { + isAsync = forStatement.isAsync(); + } + + final String iterName = isAsync ? PyNames.AITER : PyNames.ITER; + + var members = type.resolveMember(iterName, source, AccessDirection.READ, myContext); + if (members != null) results.addAll(members); + + if (results.isEmpty()) { + members = type.resolveMember(PyNames.GETITEM, source, AccessDirection.READ, myContext); + if (members != null) results.addAll(members); + } + + return RatedResolveResult.sorted(results).toArray(ResolveResult.EMPTY_ARRAY); + } + return ResolveResult.EMPTY_ARRAY; + } +} diff --git a/python/testData/resolve/keywords/forPart/aiterResolve.py b/python/testData/resolve/keywords/forPart/aiterResolve.py new file mode 100644 index 000000000000..e7cd0d55eafa --- /dev/null +++ b/python/testData/resolve/keywords/forPart/aiterResolve.py @@ -0,0 +1,10 @@ +class AsyncIterable: + async def __aiter__(self): + async def gen(): + yield 1 + return gen() + +async def run(): + async for x in AsyncIterable(): + # + pass diff --git a/python/testData/resolve/keywords/forPart/builtinListResolve.py b/python/testData/resolve/keywords/forPart/builtinListResolve.py new file mode 100644 index 000000000000..a459a11acbab --- /dev/null +++ b/python/testData/resolve/keywords/forPart/builtinListResolve.py @@ -0,0 +1,3 @@ +for x in [1, 2, 3]: + # + pass diff --git a/python/testData/resolve/keywords/forPart/onlyGetItemUnresolved.py b/python/testData/resolve/keywords/forPart/onlyGetItemUnresolved.py new file mode 100644 index 000000000000..6c38b8e7ab0a --- /dev/null +++ b/python/testData/resolve/keywords/forPart/onlyGetItemUnresolved.py @@ -0,0 +1,8 @@ +class InfiniteZeroSequence: + def __getitem__(self, i): + return 0 + + +for x in InfiniteZeroSequence(): + # + pass diff --git a/python/testData/resolve/keywords/forPart/positiveResolve.py b/python/testData/resolve/keywords/forPart/positiveResolve.py new file mode 100644 index 000000000000..49d3ee1242de --- /dev/null +++ b/python/testData/resolve/keywords/forPart/positiveResolve.py @@ -0,0 +1,7 @@ +class SingleElementIterable: + def __iter__(self): + return iter([1]) + +for x in SingleElementIterable(): + # + pass diff --git a/python/testData/resolve/keywords/forPart/unknownUnresolved.py b/python/testData/resolve/keywords/forPart/unknownUnresolved.py new file mode 100644 index 000000000000..a765b1cc7aca --- /dev/null +++ b/python/testData/resolve/keywords/forPart/unknownUnresolved.py @@ -0,0 +1,3 @@ +for x in unknown: + # + pass diff --git a/python/testSrc/com/jetbrains/python/PyForPartInKeywordResolveTest.java b/python/testSrc/com/jetbrains/python/PyForPartInKeywordResolveTest.java new file mode 100644 index 000000000000..88663b7e19ff --- /dev/null +++ b/python/testSrc/com/jetbrains/python/PyForPartInKeywordResolveTest.java @@ -0,0 +1,55 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python; + +import com.intellij.psi.PsiElement; +import com.jetbrains.python.fixtures.PyResolveTestCase; +import com.jetbrains.python.psi.PyFunction; +import com.intellij.testFramework.LightProjectDescriptor; +import com.jetbrains.python.fixtures.PyLightProjectDescriptor; +import com.jetbrains.python.psi.LanguageLevel; +import com.intellij.psi.PsiReference; + +public class PyForPartInKeywordResolveTest extends PyResolveTestCase { + + @Override + protected String getTestDataPath() { + return PythonTestUtil.getTestDataPath() + "/resolve/keywords/forPart"; + } + + @Override + protected LightProjectDescriptor getProjectDescriptor() { + return new PyLightProjectDescriptor(LanguageLevel.getLatest()); + } + + @Override + protected PsiElement doResolve() { + String fileBase = getTestName(true); + myFixture.configureByFile(fileBase + ".py"); + return PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve(); + } + + + public void testPositiveResolve() { + assertResolvesTo(PyFunction.class, "__iter__"); + PsiReference ref = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()); + String fullText = ref.getElement().getText(); + String rangeText = ref.getRangeInElement().substring(fullText); + assertEquals("in", rangeText); + } + + public void testBuiltinListResolve() { + assertResolvesTo(PyFunction.class, "__iter__"); + } + + public void testAiterResolve() { + assertResolvesTo(PyFunction.class, "__aiter__"); + } + + public void testOnlyGetItemUnresolved() { + assertResolvesTo(PyFunction.class, "__getitem__"); + } + + public void testUnknownUnresolved() { + assertUnresolved(); + } +}