PY-46668 Impl multiple arguments completion contributor

(cherry picked from commit 42efe6a20f52cfb3bd827a2b68d1fb6cc3949626)

IJ-MR-5791

GitOrigin-RevId: e2a40572b0de4919a5049ff06e437b34d78c296c

python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml

@@ -109,6 +109,8 @@
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyModulePackageCompletionContributor"/>
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyFStringLikeCompletionContributor"/>
+    <completion.contributor language="Python"
+                            implementationClass="com.jetbrains.python.codeInsight.completion.PyMultipleArgumentsCompletionContributor"/>
 
     <lang.tokenSeparatorGenerator language="Python" implementationClass="com.jetbrains.python.PyTokenSeparatorGenerator"/>
 

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyMultipleArgumentsCompletionContributor.kt

@@ -0,0 +1,129 @@
+package com.jetbrains.python.codeInsight.completion
+
+import com.intellij.codeInsight.completion.*
+import com.intellij.codeInsight.completion.util.ParenthesesInsertHandler
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.openapi.util.Key
+import com.intellij.patterns.PlatformPatterns
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.PlatformIcons
+import com.intellij.util.ProcessingContext
+import com.intellij.util.containers.ContainerUtil
+import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
+import com.jetbrains.python.extensions.inArgumentList
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.impl.PyPsiUtils
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.PyCallableParameter
+
+class PyMultipleArgumentsCompletionContributor: CompletionContributor() {
+  init {
+    extend(CompletionType.BASIC, PlatformPatterns.psiElement().inArgumentList(), MyCompletionProvider)
+  }
+
+  private object MyCompletionProvider : CompletionProvider<CompletionParameters>() {
+    override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
+      val position = parameters.position
+      val argumentIndex = getArgumentIndex(position) ?: return
+
+      val call = PsiTreeUtil.getParentOfType(position, PyCallExpression::class.java) ?: return
+      val typeEvalContext = parameters.getTypeEvalContext()
+      val resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(typeEvalContext)
+      val callableTypes = call.multiResolveCallee(resolveContext)
+      if (callableTypes.isEmpty()) return
+
+      val scopeOwner = ScopeUtil.getScopeOwner(position) ?: return
+      val names = collectNames(scopeOwner, position)
+
+      callableTypes.forEach { callableType ->
+        val callableParameters = callableType.getParameters(typeEvalContext)
+        if (callableParameters == null || callableParameters.any { it.isKeywordContainer || it.isPositionalContainer }) return@forEach
+
+        val unfilledParameters = ContainerUtil.subList(callableParameters, argumentIndex + callableType.implicitOffset)
+        val variables = collectVariablesToComplete(unfilledParameters, names)
+        if (variables.size > 1) {
+          result.addElement(createParametersLookupElement(variables, call))
+        }
+      }
+    }
+  }
+
+  companion object {
+    val MULTIPLE_ARGUMENTS_VARIANT_KEY: Key<Boolean> = Key.create("py.multiple.arguments.completion.variant")
+
+    private fun getArgumentIndex(position: PsiElement): Int? {
+      val argumentList = PsiTreeUtil.getParentOfType(position, PyArgumentList::class.java) ?: return null
+      if (argumentList.arguments.any { it is PyKeywordArgument || it is PyStarArgument }) return null
+      if (!PsiTreeUtil.isAncestor(argumentList.arguments.last(), position, false)) return null
+      return argumentList.arguments.size - 1
+    }
+
+    private fun createParametersLookupElement(variables: List<String>, call: PyCallExpression): LookupElement {
+      return LookupElementBuilder.create(variables.joinToString(", "))
+        .withIcon(PlatformIcons.VARIABLE_ICON)
+        .withInsertHandler(PyMultipleArgumentsInsertHandler(call))
+        .apply {
+          putUserData(MULTIPLE_ARGUMENTS_VARIANT_KEY, true)
+        }
+    }
+
+    private fun collectVariablesToComplete(parameters: List<PyCallableParameter>, argumentsNames: Set<String>): List<String> {
+      val variables = mutableListOf<String>()
+      var keywordsOnlyFlag = false
+
+      for (parameter in parameters) {
+        if (parameter.parameter is PySlashParameter) continue
+        if (parameter.parameter is PySingleStarParameter) {
+          keywordsOnlyFlag = true
+          continue
+        }
+
+        val paramName = parameter.name ?: return emptyList()
+        if (paramName in argumentsNames) {
+          if (!keywordsOnlyFlag) {
+            variables.add(paramName)
+          }
+          else {
+            variables.add("$paramName=$paramName")
+          }
+        }
+        else {
+          if (!parameter.hasDefaultValue()) return emptyList()
+        }
+      }
+
+      return variables
+    }
+
+    private fun collectNames(scope: ScopeOwner, position: PsiElement): Set<String> =
+      ControlFlowCache.getScope(scope).namedElements
+        .asSequence()
+        .filter { element ->
+          PsiTreeUtil.getParentOfType(element, PyListCompExpression::class.java) ?.let { listComp ->
+            PsiTreeUtil.isAncestor(listComp.resultExpression, position, false)
+          } ?: PyPsiUtils.isBefore(element, position)
+        }
+        .mapNotNull { it.name }
+        .toSet()
+  }
+}
+
+class PyMultipleArgumentsInsertHandler(private val call: PyCallExpression): ParenthesesInsertHandler<LookupElement>() {
+  override fun placeCaretInsideParentheses(context: InsertionContext?, item: LookupElement?): Boolean = false
+
+  override fun handleInsert(context: InsertionContext, item: LookupElement) {
+    val editor = context.editor
+    context.commitDocument()
+    if (call.argumentList?.closingParen == null) {
+      editor.document.insertString(context.tailOffset, ")")
+      editor.caretModel.moveToOffset(context.tailOffset)
+    }
+    else {
+      editor.caretModel.moveToOffset(context.tailOffset + 1)
+    }
+  }
+}

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PythonCompletionWeigher.java

@@ -37,6 +37,7 @@ import org.jetbrains.annotations.NotNull;
 public class PythonCompletionWeigher extends CompletionWeigher {
 
   public static final int PRIORITY_WEIGHT = 5;
+  public static final int WEIGHT_FOR_MULTIPLE_ARGUMENTS = 5;
   public static final int WEIGHT_FOR_KEYWORDS = 0;
   private static final Logger LOG = Logger.getInstance(PythonCompletionWeigher.class);
   public static final String COLLECTION_KEY = "dict key";
@@ -71,6 +72,10 @@ public class PythonCompletionWeigher extends CompletionWeigher {
       return WEIGHT_FOR_KEYWORDS;
     }
 
+    if (element.getUserData(PyMultipleArgumentsCompletionContributor.Companion.getMULTIPLE_ARGUMENTS_VARIANT_KEY()) != null) {
+      return WEIGHT_FOR_MULTIPLE_ARGUMENTS;
+    }
+
     return PyCompletionUtilsKt.FALLBACK_WEIGHT;
   }
 }

python/python-psi-impl/src/com/jetbrains/python/extensions/CaptureExt.kt

@@ -5,6 +5,7 @@ import com.intellij.patterns.PlatformPatterns.psiElement
 import com.intellij.patterns.PsiElementPattern
 import com.intellij.psi.PsiElement
 import com.jetbrains.python.PythonLanguage
+import com.jetbrains.python.psi.PyArgumentList
 import com.jetbrains.python.psi.PyClass
 import com.jetbrains.python.psi.PyFunction
 import com.jetbrains.python.psi.PyParameterList
@@ -28,3 +29,6 @@ fun PsiElementPattern.Capture<*>.afterDefInFunction(): PsiElementPattern.Capture
 
 fun PsiElementPattern.Capture<*>.inParameterList(): PsiElementPattern.Capture<out PsiElement> =
   withLanguage(PythonLanguage.getInstance()).and(psiElement().inside(PyParameterList::class.java))
+
+fun PsiElementPattern.Capture<*>.inArgumentList(): PsiElementPattern.Capture<out PsiElement> =
+  withLanguage(PythonLanguage.getInstance()).and(psiElement().inside(PyArgumentList::class.java))

python/src/com/jetbrains/python/codeInsight/mlcompletion/PyElementFeatureProvider.kt

@@ -6,6 +6,7 @@ import com.intellij.codeInsight.completion.ml.ContextFeatures
 import com.intellij.codeInsight.completion.ml.ElementFeatureProvider
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.codeInsight.lookup.LookupElement
+import com.jetbrains.python.codeInsight.completion.PyMultipleArgumentsCompletionContributor
 import com.jetbrains.python.codeInsight.mlcompletion.prev2calls.PyPrevCallsCompletionFeatures
 import com.jetbrains.python.psi.PyParameter
 
@@ -100,6 +101,10 @@ class PyElementFeatureProvider : ElementFeatureProvider {
       }}
     }
 
+    element.getUserData(PyMultipleArgumentsCompletionContributor.MULTIPLE_ARGUMENTS_VARIANT_KEY)?.let {
+      result["is_multiple_arguments"] = MLFeatureValue.binary(true)
+    }
+
     return result
   }
 }

python/testData/multipleArgumentsCompletion/argumentParameter.after.py

@@ -0,0 +1,7 @@
+def foo(bar, baz):
+    pass
+
+
+def main(bar):
+    for baz in []:
+        foo(bar, baz)<caret>

python/testData/multipleArgumentsCompletion/argumentParameter.py

@@ -0,0 +1,7 @@
+def foo(bar, baz):
+    pass
+
+
+def main(bar):
+    for baz in []:
+        foo(<caret>)

python/testData/multipleArgumentsCompletion/classConstructor.after.py

@@ -0,0 +1,12 @@
+class Foo:
+    def __init__(self, x, y, z):
+        self.x = x
+        self.y = y
+        self.z = z
+
+
+def foo():
+    x = 1
+    y = 2
+    z = 3
+    return Foo(x, y, z)<caret>

python/testData/multipleArgumentsCompletion/classConstructor.py

@@ -0,0 +1,12 @@
+class Foo:
+    def __init__(self, x, y, z):
+        self.x = x
+        self.y = y
+        self.z = z
+
+
+def foo():
+    x = 1
+    y = 2
+    z = 3
+    return Foo(<caret>)

python/testData/multipleArgumentsCompletion/classMember.after.py

@@ -0,0 +1,9 @@
+class Foo:
+    def foo(self, x, y, z):
+        pass
+
+    def bar(self):
+        x = 1
+        y = 2
+        z = 3
+        self.foo(x, y, z)<caret>

python/testData/multipleArgumentsCompletion/classMember.py

@@ -0,0 +1,9 @@
+class Foo:
+    def foo(self, x, y, z):
+        pass
+
+    def bar(self):
+        x = 1
+        y = 2
+        z = 3
+        self.foo(<caret>)

python/testData/multipleArgumentsCompletion/functionWithThreeArgsTwoSuggested.after.py

@@ -0,0 +1,8 @@
+def foo(x, y, z):
+    pass
+
+
+z = 22
+x = 42
+y = 100500
+foo(x, y, z)<caret>

python/testData/multipleArgumentsCompletion/functionWithThreeArgsTwoSuggested.py

@@ -0,0 +1,8 @@
+def foo(x, y, z):
+    pass
+
+
+z = 22
+x = 42
+y = 100500
+foo(x, <caret>

python/testData/multipleArgumentsCompletion/functionWithTwoArgs.after.py

@@ -0,0 +1,7 @@
+def foo(x, y):
+    pass
+
+
+x = 42
+y = 100500
+foo(x, y)<caret>

python/testData/multipleArgumentsCompletion/functionWithTwoArgs.py

@@ -0,0 +1,7 @@
+def foo(x, y):
+    pass
+
+
+x = 42
+y = 100500
+foo(<caret>)

python/testData/multipleArgumentsCompletion/listComprehension.after.py

@@ -0,0 +1,6 @@
+def foo(bar, baz):
+    pass
+
+
+def main():
+    return [foo(bar, baz)<caret> for bar, baz in []]

python/testData/multipleArgumentsCompletion/listComprehension.py

@@ -0,0 +1,6 @@
+def foo(bar, baz):
+    pass
+
+
+def main():
+    return [foo(<caret>) for bar, baz in []]

python/testData/multipleArgumentsCompletion/notSuggestIfHaveArgumentsRight.py

@@ -0,0 +1,8 @@
+def foo(x, y, z):
+    pass
+
+
+x = 42
+y = 100500
+z = 42
+foo(x, <caret>, z)

python/testData/multipleArgumentsCompletion/notSuggestIfNotEnoughArguments.py

@@ -0,0 +1,7 @@
+def foo(x, y, z):
+    pass
+
+
+x = 42
+y = 100500
+foo(<caret>)

python/testData/multipleArgumentsCompletion/notSuggestIfNotEnoughArgumentsBeforeCaret.py

@@ -0,0 +1,8 @@
+def foo(x, y, z):
+    pass
+
+
+x = 42
+y = 100500
+foo(<caret>)
+z = 33

python/testData/multipleArgumentsCompletion/notSuggestIfNotEnoughArgumentsInTheScopeOfFunction.py

@@ -0,0 +1,9 @@
+def foo(x, y, z):
+    pass
+
+
+z = 33
+def bar():
+    x = 42
+    y = 100500
+    foo(<caret>)

python/testData/multipleArgumentsCompletion/notSuggestIfNotEnoughArgumentsInTheScopeOfNestedFunction.py

@@ -0,0 +1,10 @@
+def foo(x, y, z):
+    pass
+
+
+def bar():
+    z = 33
+    def baz():
+        x = 42
+        y = 100500
+        foo(<caret>)

python/testData/multipleArgumentsCompletion/notSuggestIfOnlyOneVariable.py

@@ -0,0 +1,5 @@
+def foo(x, y=42)
+    pass
+
+x = 42
+foo(<caret>)

python/testData/multipleArgumentsCompletion/notSuggestIfTargetOutOfScopeFunction.py

@@ -0,0 +1,7 @@
+def foo(x, y, z):
+    z = 42
+
+
+x = 42
+y = 100500
+foo(<caret>)

python/testData/multipleArgumentsCompletion/notSuggestInKeywordArgument.py

@@ -0,0 +1,8 @@
+def foo(x, y, z):
+    pass
+
+
+x = 42
+y = 100500
+z = 42
+foo(x=<caret>

python/testData/multipleArgumentsCompletion/notSuggestKeywordContainer.py

@@ -0,0 +1,7 @@
+def foo(x, **y):
+    pass
+
+
+x = 42
+y = 22
+foo(<caret>

python/testData/multipleArgumentsCompletion/notSuggestPositionalContainer.py

@@ -0,0 +1,7 @@
+def foo(x, *y):
+    pass
+
+
+x = 42
+y = 22
+foo(<caret>

python/testData/multipleArgumentsCompletion/overloads.py

@@ -0,0 +1,24 @@
+from typing import overload, Any
+
+
+@overload
+def bar(a: int, b: int) -> None:
+    ...
+
+
+@overload
+def bar(c: str, d: str) -> None:
+    ...
+
+
+def bar(*args: Any, **kwargs: Any) -> None:
+    ...
+
+
+def foo():
+    a = 1
+    b = 2
+    c = 3
+    d = 4
+
+    bar(<caret>)

python/testData/multipleArgumentsCompletion/singleStarParameter.after.py

@@ -0,0 +1,8 @@
+def foo(a, *, b):
+    print(a, b)
+
+
+def egg():
+    a = 1
+    b = 2
+    foo(a, b=b)<caret>

python/testData/multipleArgumentsCompletion/singleStarParameter.py

@@ -0,0 +1,8 @@
+def foo(a, *, b):
+    print(a, b)
+
+
+def egg():
+    a = 1
+    b = 2
+    foo(<caret>)

python/testData/multipleArgumentsCompletion/skipArgumentsForParametersWithDefaultValue.after.py

@@ -0,0 +1,7 @@
+def foo(x, y, z=42):
+    pass
+
+
+x = 42
+y = 100500
+foo(x, y)<caret>

python/testData/multipleArgumentsCompletion/skipArgumentsForParametersWithDefaultValue.py

@@ -0,0 +1,7 @@
+def foo(x, y, z=42):
+    pass
+
+
+x = 42
+y = 100500
+foo(<caret>

python/testData/multipleArgumentsCompletion/skipDefaultValueIfNotFromFunctionContext.after.py

@@ -0,0 +1,9 @@
+def foo(x, y, z=42):
+    pass
+
+
+z = 42
+def bar():
+  x = 42
+  y = 100500
+  foo(x, y)<caret>

python/testData/multipleArgumentsCompletion/skipDefaultValueIfNotFromFunctionContext.py

@@ -0,0 +1,9 @@
+def foo(x, y, z=42):
+    pass
+
+
+z = 42
+def bar():
+  x = 42
+  y = 100500
+  foo(<caret>

python/testData/multipleArgumentsCompletion/slashAndSingleStarParameter.after.py

@@ -0,0 +1,9 @@
+def foo(a, /, b, *, c):
+    print(a, b, c)
+
+
+def egg():
+    a = 1
+    b = 2
+    c = 3
+    foo(a, b, c=c)<caret>

python/testData/multipleArgumentsCompletion/slashAndSingleStarParameter.py

@@ -0,0 +1,9 @@
+def foo(a, /, b, *, c):
+    print(a, b, c)
+
+
+def egg():
+    a = 1
+    b = 2
+    c = 3
+    foo(<caret>

python/testData/multipleArgumentsCompletion/slashParameter.after.py

@@ -0,0 +1,8 @@
+def foo(a, /, b):
+    print(a, b)
+
+
+def egg():
+    a = 1
+    b = 2
+    foo(a, b)<caret>

python/testData/multipleArgumentsCompletion/slashParameter.py

@@ -0,0 +1,8 @@
+def foo(a, /, b):
+    print(a, b)
+
+
+def egg():
+    a = 1
+    b = 2
+    foo(<caret>

python/testData/multipleArgumentsCompletion/suggestArgumentsForParametersWithDefaultValue.after.py

@@ -0,0 +1,8 @@
+def foo(x, y, z=42):
+    pass
+
+
+x = 42
+y = 100500
+z = 42
+foo(x, y, z)<caret>

python/testData/multipleArgumentsCompletion/suggestArgumentsForParametersWithDefaultValue.py

@@ -0,0 +1,8 @@
+def foo(x, y, z=42):
+    pass
+
+
+x = 42
+y = 100500
+z = 42
+foo(<caret>

python/testSrc/com/jetbrains/python/PyMultipleArgumentsCompletionTest.kt

@@ -0,0 +1,117 @@
+// Copyright 2000-2021 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
+
+import com.intellij.codeInsight.lookup.Lookup
+import com.jetbrains.python.fixtures.PyTestCase
+
+class PyMultipleArgumentsCompletionTest: PyTestCase() {
+  fun testFunctionWithTwoArgs() {
+    doTestVariantsContainFinished("x, y")
+  }
+
+  fun testFunctionWithThreeArgsTwoSuggested() {
+    doTestVariantsContainFinished("y, z")
+  }
+
+  fun testSuggestArgumentsForParametersWithDefaultValue() {
+    doTestVariantsContainFinished("x, y, z")
+  }
+
+  fun testSkipArgumentsForParametersWithDefaultValue() {
+    doTestVariantsContainFinished("x, y")
+  }
+
+  fun testSkipDefaultValueIfNotFromFunctionContext() {
+    doTestVariantsContainFinished("x, y")
+  }
+
+  fun testClassConstructor() {
+    doTestVariantsContainFinished("x, y, z")
+  }
+
+  fun testClassMember() {
+    doTestVariantsContainFinished("x, y, z")
+  }
+
+  fun testArgumentParameter() {
+    doTestVariantsContainFinished("bar, baz")
+  }
+
+  fun testListComprehension() {
+    doTestVariantsContainFinished("bar, baz")
+  }
+
+  fun testOverloads() {
+    doTestVariantsContain("a, b", "c, d")
+  }
+
+  fun testSingleStarParameter() {
+    doTestVariantsContainFinished("a, b=b")
+  }
+
+  fun testSlashParameter() {
+    doTestVariantsContainFinished("a, b")
+  }
+
+  fun testSlashAndSingleStarParameter() {
+    doTestVariantsContainFinished("a, b, c=c")
+  }
+
+  fun testNotSuggestIfNotEnoughArguments() {
+    doTestVariantsNotContain("x, y", "x, y, z")
+  }
+
+  fun testNotSuggestIfNotEnoughArgumentsBeforeCaret() {
+    doTestVariantsNotContain("x, y", "x, y, z")
+  }
+
+  fun testNotSuggestIfNotEnoughArgumentsInTheScopeOfFunction() {
+    doTestVariantsNotContain("x, y", "x, y, z")
+  }
+
+  fun testNotSuggestIfNotEnoughArgumentsInTheScopeOfNestedFunction() {
+    doTestVariantsNotContain("x, y", "x, y, z")
+  }
+
+  fun testNotSuggestIfTargetOutOfScopeFunction() {
+    doTestVariantsNotContain("x, y", "x, y, z")
+  }
+
+  fun testNotSuggestIfHaveArgumentsRight() {
+    doTestVariantsNotContain("y, z")
+  }
+
+  fun testNotSuggestInKeywordArgument() {
+    doTestVariantsNotContain("x, y, z", "y, z")
+  }
+
+  fun testNotSuggestKeywordContainer() {
+    doTestVariantsNotContain("x, y")
+  }
+
+  fun testNotSuggestPositionalContainer() {
+    doTestVariantsNotContain("x, y")
+  }
+
+  fun testNotSuggestIfOnlyOneVariable() {
+    assertEquals(1, doTestByTestName().count { it == "x" })
+  }
+
+  private fun doTestByTestName(): List<String?> {
+    val testName = "multipleArgumentsCompletion/${getTestName(true)}"
+    myFixture.configureByFile("$testName.py")
+    myFixture.completeBasic()
+    return myFixture.lookupElementStrings!!
+  }
+
+  private fun doTestVariantsContain(vararg elements: String) = assertContainsElements(doTestByTestName(), *elements)
+
+  private fun doTestVariantsContainFinished(element: String) {
+    doTestVariantsContain(element)
+    myFixture.lookup.currentItem = myFixture.lookupElements.find { it.lookupString == element }
+    myFixture.finishLookup(Lookup.NORMAL_SELECT_CHAR)
+    myFixture.checkResultByFile("multipleArgumentsCompletion/${getTestName(true)}.after.py")
+  }
+
+  private fun doTestVariantsNotContain(vararg elements: String) = assertDoesntContain(doTestByTestName(), *elements)
+}