PY-82412 support go to definition on loop constructs [python]

Merge-request: IJ-MR-181460
Merged-by: Tatiana Ber <tatiana.ber@jetbrains.com>

GitOrigin-RevId: 81affdaa7f2e8747557444eb732f0fe106e0d224

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);
+  }
 }

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<PyElement> {
+
+  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<RatedResolveResult> 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;
+  }
+}

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():
+                # <ref>
+        pass

python/testData/resolve/keywords/forPart/builtinListResolve.py

@@ -0,0 +1,3 @@
+for x in [1, 2, 3]:
+    # <ref>
+    pass

python/testData/resolve/keywords/forPart/onlyGetItemUnresolved.py

@@ -0,0 +1,8 @@
+class InfiniteZeroSequence:
+    def __getitem__(self, i):
+        return 0
+
+
+for x in InfiniteZeroSequence():
+    # <ref>
+    pass

python/testData/resolve/keywords/forPart/positiveResolve.py

@@ -0,0 +1,7 @@
+class SingleElementIterable:
+    def __iter__(self):
+        return iter([1])
+
+for x in SingleElementIterable():
+    # <ref>
+    pass

python/testData/resolve/keywords/forPart/unknownUnresolved.py

@@ -0,0 +1,3 @@
+for x in unknown:
+    # <ref>
+    pass

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();
+  }
+}