PY-40480 Completion for argument passed to parameter hinted as Literal

GitOrigin-RevId: 462b03dc35eef70bc873b568cae239d40f1d4974

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

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

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

@@ -0,0 +1,65 @@
+package com.jetbrains.python.codeInsight.completion
+
+import com.intellij.codeInsight.completion.*
+import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.patterns.PlatformPatterns.psiElement
+import com.intellij.patterns.StandardPatterns.or
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.ui.IconManager
+import com.intellij.ui.PlatformIcons
+import com.intellij.util.ProcessingContext
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.impl.PyCallExpressionHelper
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.PyLiteralType
+import com.jetbrains.python.psi.types.PyTypeUtil
+import com.jetbrains.python.psi.types.TypeEvalContext
+
+class PyLiteralTypeCompletionContributor : CompletionContributor() {
+  init {
+    extend(
+      CompletionType.BASIC,
+      or(
+        psiElement().withSuperParent(2, PyKeywordArgument::class.java),
+        psiElement().withSuperParent(2, PyArgumentList::class.java),
+        psiElement().withSuperParent(2, PySubscriptionExpression::class.java)
+      ),
+      PyLiteralTypeCompletionProvider()
+    )
+  }
+}
+
+private class PyLiteralTypeCompletionProvider : CompletionProvider<CompletionParameters?>() {
+  override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
+    val position = parameters.position.parent as? PyExpression ?: return
+    val callSiteExpr = PsiTreeUtil.getParentOfType(position, PyCallSiteExpression::class.java) ?: return
+
+    val parent = position.parent
+    val argumentExpr = if (parent is PyKeywordArgument && parent.valueExpression == position) parent else position
+    val typeEvalContext = TypeEvalContext.codeCompletion(position.project, position.containingFile)
+    val types = PyCallExpressionHelper.mapArguments(callSiteExpr, PyResolveContext.defaultContext(typeEvalContext))
+      .mapNotNull { it.mappedParameters[argumentExpr]?.getArgumentType(typeEvalContext) }
+      .flatMap { PyTypeUtil.toStream(it) }
+      .filterIsInstance<PyLiteralType>()
+
+    for (type in types) {
+      val expression = type.expression
+      if (position is PyStringLiteralExpression) {
+        if (expression is PyStringLiteralExpression) {
+          addToResult(result, expression.stringValue)
+        }
+      }
+      else if (expression is PyStringLiteralExpression || expression is PyNumericLiteralExpression) {
+        addToResult(result, expression.text)
+      }
+    }
+  }
+
+  private fun addToResult(result: CompletionResultSet, lookupString: String) {
+    result.addElement(
+      LookupElementBuilder
+        .create(lookupString)
+        .withIcon(IconManager.getInstance().getPlatformIcon(PlatformIcons.Parameter))
+    )
+  }
+}

python/testData/completion/literalType/inCallExpression.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def f(x: Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]) -> None:
+    pass
+
+
+f(<caret>)

python/testData/completion/literalType/inDoubleQuotedString.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def f(x: Literal["foo", 3, "bar"]) -> None:
+    pass
+
+
+f("<caret>")

python/testData/completion/literalType/inKeywordArgument.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def f(x: Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]) -> None:
+    pass
+
+
+f(x=<caret>)

python/testData/completion/literalType/inSingleQuotedString.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def f(x: Literal["foo", 3, "bar"]) -> None:
+    pass
+
+
+f('<caret>')

python/testData/completion/literalType/inSubscriptionExpression.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+class A:
+    def __getitem__(self, item: Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]) -> str:
+        pass
+
+
+A()[<caret>]

python/testData/completion/literalType/kvarargs.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def foo(**x: Literal[Literal[Literal["aa", 'bbb'], "zzz"], 5]):
+    pass
+
+
+foo(key0="<caret>")

python/testData/completion/literalType/literalWithDoubleQuotedStringParameter.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def f(x: Literal["yyy"]) -> None:
+    pass
+
+
+f(<caret>)

python/testData/completion/literalType/literalWithSingleParameter.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def f(x: Literal[22]) -> None:
+    pass
+
+
+f(<caret>)

python/testData/completion/literalType/literalWithSingleQuotedStringParameter.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def f(x: Literal['xxx']) -> None:
+    pass
+
+
+f(<caret>)

python/testData/completion/literalType/nestedArgumentLists.py

@@ -0,0 +1,12 @@
+from typing import Literal
+
+
+def expects_literal(x: Literal["foo", "bar"]) -> None:
+    pass
+
+
+def f(x: str) -> str:
+    pass
+
+
+expects_literal(f("<caret>"))

python/testData/completion/literalType/varargs.py

@@ -0,0 +1,8 @@
+from typing import Literal
+
+
+def foo(*x: Literal[Literal[Literal["aa", 'bbb'], "zzz"], 5]):
+    pass
+
+
+foo("<caret>")

python/testSrc/com/jetbrains/python/PyLiteralTypeCompletionTest.kt

@@ -0,0 +1,63 @@
+// Copyright 2000-2024 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.testFramework.TestDataPath
+import com.jetbrains.python.fixtures.PyTestCase
+
+// PY-40480
+@TestDataPath("\$CONTENT_ROOT/../testData/completion/literalType")
+class PyLiteralTypeCompletionTest : PyTestCase() {
+  fun testInCallExpression() {
+    doTestCompletionVariantsContains("inCallExpression.py", "1", "2", "3", "\"foo\"", "5", "None")
+  }
+
+  fun testInKeywordArgument() {
+    doTestCompletionVariantsContains("inKeywordArgument.py", "1", "2", "3", "\"foo\"", "5", "None")
+  }
+
+  fun testInSubscriptionExpression() {
+    doTestCompletionVariantsContains("inSubscriptionExpression.py", "1", "2", "3", "\"foo\"", "5", "None")
+  }
+
+  fun testInDoubleQuotedString() {
+    myFixture.testCompletionVariants("inDoubleQuotedString.py", "bar", "foo")
+  }
+
+  fun testInSingleQuotedString() {
+    myFixture.testCompletionVariants("inSingleQuotedString.py", "bar", "foo")
+  }
+
+  fun testLiteralWithSingleParameter() {
+    doTestCompletionVariantsContains("literalWithSingleParameter.py", "22")
+  }
+
+  fun testLiteralWithDoubleQuotedStringParameter() {
+    doTestCompletionVariantsContains("literalWithDoubleQuotedStringParameter.py", "\"yyy\"")
+  }
+
+  fun testLiteralWithSingleQuotedStringParameter() {
+    doTestCompletionVariantsContains("literalWithSingleQuotedStringParameter.py", "'xxx'")
+  }
+
+  fun testVarargs() {
+    myFixture.testCompletionVariants("varargs.py", "aa", "bbb", "zzz")
+  }
+
+  fun testKVarargs() {
+    myFixture.testCompletionVariants("kvarargs.py", "aa", "bbb", "zzz")
+  }
+
+  fun testNestedArgumentLists() {
+    myFixture.testCompletionVariants("nestedArgumentLists.py")
+  }
+
+  override fun getTestDataPath(): String {
+    return super.getTestDataPath() + "/completion/literalType"
+  }
+
+  private fun doTestCompletionVariantsContains(fileBefore: String, vararg items: String) {
+    val result = myFixture.getCompletionVariants(fileBefore)
+    assertNotNull(result)
+    assertContainsElements(result!!, *items)
+  }
+}