PY-32269: Fixed bug with completion of parameters before = sign.

GitOrigin-RevId: bb619c210c66a595ff015494b4b553a15cf24c1e

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -1139,13 +1139,17 @@ public class PyUtil {
    * @return lookup element
    */
   @NotNull
-  public static LookupElement createNamedParameterLookup(@NotNull String name, @NotNull PsiFile settingsAnchor) {
+  public static LookupElement createNamedParameterLookup(@NotNull String name, @NotNull PsiFile settingsAnchor, boolean addSuffix) {
     final String suffix;
-    if (CodeStyle.getCustomSettings(settingsAnchor, PyCodeStyleSettings.class).SPACE_AROUND_EQ_IN_KEYWORD_ARGUMENT) {
-      suffix = " = ";
-    }
-    else {
-      suffix = "=";
+    if (addSuffix) {
+      if (CodeStyle.getCustomSettings(settingsAnchor, PyCodeStyleSettings.class).SPACE_AROUND_EQ_IN_KEYWORD_ARGUMENT) {
+        suffix = " = ";
+      }
+      else {
+        suffix = "=";
+      }
+    } else {
+      suffix = "";
     }
     LookupElementBuilder lookupElementBuilder = LookupElementBuilder.create(name + suffix).withIcon(PlatformIcons.PARAMETER_ICON);
     lookupElementBuilder = lookupElementBuilder.withInsertHandler(OverwriteEqualsInsertHandler.INSTANCE);

python/src/com/jetbrains/python/psi/impl/PyKeywordArgumentReference.java

@@ -15,10 +15,15 @@
  */
 package com.jetbrains.python.psi.impl;
 
+import com.google.common.collect.Lists;
+import com.intellij.codeInsight.completion.CompletionUtil;
+import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.references.KeywordArgumentCompletionUtil;
+import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -34,6 +39,19 @@ public class PyKeywordArgumentReference extends PsiReferenceBase.Poly<PyKeywordA
     super(element, textRange, true);
   }
 
+  public Object[] getVariants() {
+    final List<LookupElement> ret = Lists.newArrayList();
+
+    final PyKeywordArgument originalElement = CompletionUtil.getOriginalElement(myElement);
+    final PyKeywordArgument element = originalElement != null ? originalElement : myElement;
+
+    KeywordArgumentCompletionUtil
+      .collectFunctionArgNames(element, ret, TypeEvalContext.codeCompletion(element.getProject(), element.getContainingFile()), false);
+
+    return ret.toArray();
+  }
+
+
   @NotNull
   @Override
   public ResolveResult[] multiResolve(boolean incompleteCode) {
@@ -47,7 +65,7 @@ public class PyKeywordArgumentReference extends PsiReferenceBase.Poly<PyKeywordA
     }
     final PyExpression callee = ((PyCallExpression)call).getCallee();
     if (callee == null) return ResolveResult.EMPTY_ARRAY;
-    final PsiPolyVariantReference calleeReference = (PsiPolyVariantReference) callee.getReference();
+    final PsiPolyVariantReference calleeReference = (PsiPolyVariantReference)callee.getReference();
     if (calleeReference == null) return ResolveResult.EMPTY_ARRAY;
     final ResolveResult[] calleeCandidates = calleeReference.multiResolve(incompleteCode);
     List<ResolveResult> resultList = new ArrayList<>();

python/src/com/jetbrains/python/psi/impl/references/KeywordArgumentCompletionUtil.java

@@ -24,7 +24,7 @@ import java.util.Set;
 import static com.jetbrains.python.psi.PyUtil.as;
 
 public class KeywordArgumentCompletionUtil {
-  public static void collectFunctionArgNames(PyElement element, List<? super LookupElement> ret, @NotNull final TypeEvalContext context) {
+  public static void collectFunctionArgNames(PyElement element, List<? super LookupElement> ret, @NotNull final TypeEvalContext context, final boolean addSuffix) {
     PyCallExpression callExpr = PsiTreeUtil.getParentOfType(element, PyCallExpression.class);
     if (callExpr != null) {
       PyExpression callee = callExpr.getCallee();
@@ -39,7 +39,7 @@ public class KeywordArgumentCompletionUtil {
         final List<LookupElement> extra = PyTypeUtil.toStream(calleeType)
           .select(PyCallableType.class)
           .flatMap(type -> collectParameterNamesFromType(type, callExpr, context).stream())
-          .map(name -> PyUtil.createNamedParameterLookup(name, element.getContainingFile()))
+          .map(name -> PyUtil.createNamedParameterLookup(name, element.getContainingFile(), addSuffix))
           .toList();
 
         ret.addAll(extra);

python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java

@@ -707,8 +707,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
 
     // This method is probably called for completion, so use appropriate context here
     // in a call, include function's arg names
-    KeywordArgumentCompletionUtil.collectFunctionArgNames(element, ret, TypeEvalContext.codeCompletion(element.getProject(), element.getContainingFile()));
-
+    KeywordArgumentCompletionUtil.collectFunctionArgNames(element, ret, TypeEvalContext.codeCompletion(element.getProject(), element.getContainingFile()), true);
     // include builtin names
     final PyFile builtinsFile = builtinCache.getBuiltinsFile();
     if (builtinsFile != null) {

python/testData/completion/paramCompletionWithEquals.after.py

@@ -0,0 +1,3 @@
+def func(myArg) : pass
+
+func(myArg=)

python/testData/completion/paramCompletionWithEquals.py

@@ -0,0 +1,3 @@
+def func(myArg) : pass
+
+func(my<caret>=)

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -1524,6 +1524,11 @@ public class PythonCompletionTest extends PyTestCase {
     assertDoesntContain(suggested, "foo");
   }
 
+  // PY-32269
+  public void testParamCompletionWithEquals() {
+    doTest();
+  }
+
 
   private void assertNoVariantsInExtendedCompletion() {
     myFixture.copyDirectoryToProject(getTestName(true), "");