Collect completion variants for doc reference in the same way as resolving it (PY-31517)

GitOrigin-RevId: 7731555e5a975d41c8ee40b44f8a2c6951a7fa39

python/python-psi-impl/src/com/jetbrains/python/documentation/doctest/PyDocReference.java

@@ -1,11 +1,12 @@
 // Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python.documentation.doctest;
 
-import com.google.common.collect.Lists;
 import com.intellij.codeInsight.completion.CompletionUtilCoreImpl;
+import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.lang.ASTNode;
 import com.intellij.lang.annotation.HighlightSeverity;
 import com.intellij.lang.injection.InjectedLanguageManager;
+import com.intellij.openapi.util.Condition;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiElement;
@@ -13,6 +14,7 @@ import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiLanguageInjectionHost;
 import com.intellij.psi.ResolveResult;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.ArrayUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
@@ -21,11 +23,12 @@ import com.jetbrains.python.psi.PyUtil.StringNodeInfo;
 import com.jetbrains.python.psi.impl.references.PyReferenceImpl;
 import com.jetbrains.python.psi.resolve.*;
 import com.jetbrains.python.psi.types.TypeEvalContext;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 
 /**
  * User : ktisha
@@ -51,7 +54,7 @@ public class PyDocReference extends PyReferenceImpl {
       if (referencedName == null) return ResolveResult.EMPTY_ARRAY;
 
       if (host != null) {
-        final List<Pair<PsiElement,TextRange>> files = languageManager.getInjectedPsiFiles(host);
+        final List<Pair<PsiElement, TextRange>> files = languageManager.getInjectedPsiFiles(host);
         if (files != null) {
           for (Pair<PsiElement, TextRange> pair : files) {
             if (pair.getFirst() instanceof PyFile) {
@@ -68,7 +71,7 @@ public class PyDocReference extends PyReferenceImpl {
           }
         }
         final PyResolveProcessor processor = new PyResolveProcessor(referencedName);
-        final ScopeOwner scopeOwner = getHostScopeOwner();
+        final ScopeOwner scopeOwner = ScopeUtil.getScopeOwner(host);
         if (scopeOwner != null) {
           final PsiFile topLevel = scopeOwner.getContainingFile();
           PyResolveUtil.scopeCrawlUp(processor, scopeOwner, referencedName, topLevel);
@@ -136,33 +139,38 @@ public class PyDocReference extends PyReferenceImpl {
   }
 
   @Override
+  @NotNull
   public Object @NotNull [] getVariants() {
-    final ArrayList<Object> ret = Lists.newArrayList(super.getVariants());
-    final PyQualifiedExpression originalElement = CompletionUtilCoreImpl.getOriginalElement(myElement);
-    final PyQualifiedExpression element = originalElement != null ? originalElement : myElement;
+    final Object[] results = super.getVariants();
 
-    final ScopeOwner scopeOwner = getHostScopeOwner();
-    if (scopeOwner != null) {
-      final CompletionVariantsProcessor processor = new CompletionVariantsProcessor(element);
-      PyResolveUtil.scopeCrawlUp(processor, scopeOwner, null, null);
-      ret.addAll(processor.getResultList());
-    }
-    return ret.toArray();
-  }
-
-
-  @Nullable
-  private ScopeOwner getHostScopeOwner() {
     final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myElement.getProject());
     final PsiLanguageInjectionHost host = languageManager.getInjectionHost(myElement);
-    if (host != null) {
-      final PsiFile file = host.getContainingFile();
-      ScopeOwner result = ScopeUtil.getScopeOwner(host);
-      if (result == null && file instanceof ScopeOwner) {
-        result = (ScopeOwner)file;
+    if (host == null) return results;
+
+    final PyQualifiedExpression originalElement = CompletionUtilCoreImpl.getOriginalElement(myElement);
+    final PyQualifiedExpression element = originalElement != null ? originalElement : myElement;
+    final CompletionVariantsProcessor processor = new CompletionVariantsProcessor(element, null, filterForPresentedNames(results));
+
+    final List<Pair<PsiElement, TextRange>> files = languageManager.getInjectedPsiFiles(host);
+    if (files != null) {
+      for (Pair<PsiElement, TextRange> pair : files) {
+        if (pair.getFirst() instanceof PyFile) {
+          PyResolveUtil.scopeCrawlUp(processor, (ScopeOwner)pair.getFirst(), null, pair.getFirst());
+        }
       }
-      return result;
     }
-    return null;
+    final ScopeOwner scopeOwner = ScopeUtil.getScopeOwner(host);
+    if (scopeOwner != null) {
+      PyResolveUtil.scopeCrawlUp(processor, scopeOwner, null, scopeOwner.getContainingFile());
+    }
+
+    return ArrayUtil.mergeArrayAndCollection(results, processor.getResultList(), Object[]::new);
+  }
+
+  @Nullable
+  private static Condition<String> filterForPresentedNames(@NotNull Object[] variants) {
+    if (variants.length == 0) return null;
+    final Set<String> seenNames = StreamEx.of(variants).select(LookupElement.class).map(LookupElement::getLookupString).toSet();
+    return s -> !seenNames.contains(s);
   }
 }

python/testData/completion/doctest/functionDeclaredInDoctestBefore.py

@@ -0,0 +1,7 @@
+"""
+>>> def foo(bar):
+...     print(bar)
+
+>>> fo<caret>
+Hello
+"""

python/testSrc/com/jetbrains/python/PythonDoctestCompletionTest.java

@@ -16,6 +16,7 @@
 package com.jetbrains.python;
 
 import com.intellij.codeInsight.lookup.LookupElement;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.fixtures.PyTestCase;
 
 /**
@@ -27,14 +28,8 @@ public class PythonDoctestCompletionTest extends PyTestCase {
     final String testName = getTestName(true);
     myFixture.configureByFile(testName + ".py");
     final LookupElement[] elements = myFixture.completeBasic();
-    if (elements != null) {
-      for (LookupElement lookup : elements) {
-        LOG.debug(lookup.getLookupString());
-        if (lookup.getLookupString().equals(expected))
-          return;
-      }
-    }
-    fail();
+    assertNotNull(elements);
+    assertContainsElements(ContainerUtil.map(elements, LookupElement::getLookupString), expected);
   }
 
   public void testForInDoctest() {
@@ -54,6 +49,11 @@ public class PythonDoctestCompletionTest extends PyTestCase {
     doDoctestTest("foo");
   }
 
+  // PY-31517
+  public void testFunctionDeclaredInDoctestBefore() {
+    doDoctestTest("foo");
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/completion/doctest";