diff --git a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml index 921e0155cb78..f51b390020c7 100644 --- a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml +++ b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml @@ -148,6 +148,8 @@ implementationClass="com.jetbrains.python.codeInsight.completion.PyFStringLikeCompletionContributor"/> + diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyLiteralTypeCompletionContributor.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyLiteralTypeCompletionContributor.kt new file mode 100644 index 000000000000..433311f028c6 --- /dev/null +++ b/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() { + 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() + + 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)) + ) + } +} \ No newline at end of file diff --git a/python/testData/completion/literalType/inCallExpression.py b/python/testData/completion/literalType/inCallExpression.py new file mode 100644 index 000000000000..ddca66b4c69c --- /dev/null +++ b/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() diff --git a/python/testData/completion/literalType/inDoubleQuotedString.py b/python/testData/completion/literalType/inDoubleQuotedString.py new file mode 100644 index 000000000000..cf32bd96c4ca --- /dev/null +++ b/python/testData/completion/literalType/inDoubleQuotedString.py @@ -0,0 +1,8 @@ +from typing import Literal + + +def f(x: Literal["foo", 3, "bar"]) -> None: + pass + + +f("") diff --git a/python/testData/completion/literalType/inKeywordArgument.py b/python/testData/completion/literalType/inKeywordArgument.py new file mode 100644 index 000000000000..3f7946ec3ccd --- /dev/null +++ b/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=) diff --git a/python/testData/completion/literalType/inSingleQuotedString.py b/python/testData/completion/literalType/inSingleQuotedString.py new file mode 100644 index 000000000000..575c464fcbe6 --- /dev/null +++ b/python/testData/completion/literalType/inSingleQuotedString.py @@ -0,0 +1,8 @@ +from typing import Literal + + +def f(x: Literal["foo", 3, "bar"]) -> None: + pass + + +f('') diff --git a/python/testData/completion/literalType/inSubscriptionExpression.py b/python/testData/completion/literalType/inSubscriptionExpression.py new file mode 100644 index 000000000000..62313d5fab0f --- /dev/null +++ b/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()[] \ No newline at end of file diff --git a/python/testData/completion/literalType/kvarargs.py b/python/testData/completion/literalType/kvarargs.py new file mode 100644 index 000000000000..6491aa1c61ce --- /dev/null +++ b/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="") diff --git a/python/testData/completion/literalType/literalWithDoubleQuotedStringParameter.py b/python/testData/completion/literalType/literalWithDoubleQuotedStringParameter.py new file mode 100644 index 000000000000..87646abd1341 --- /dev/null +++ b/python/testData/completion/literalType/literalWithDoubleQuotedStringParameter.py @@ -0,0 +1,8 @@ +from typing import Literal + + +def f(x: Literal["yyy"]) -> None: + pass + + +f() diff --git a/python/testData/completion/literalType/literalWithSingleParameter.py b/python/testData/completion/literalType/literalWithSingleParameter.py new file mode 100644 index 000000000000..01e00bdff6a9 --- /dev/null +++ b/python/testData/completion/literalType/literalWithSingleParameter.py @@ -0,0 +1,8 @@ +from typing import Literal + + +def f(x: Literal[22]) -> None: + pass + + +f() diff --git a/python/testData/completion/literalType/literalWithSingleQuotedStringParameter.py b/python/testData/completion/literalType/literalWithSingleQuotedStringParameter.py new file mode 100644 index 000000000000..a8295555fe6e --- /dev/null +++ b/python/testData/completion/literalType/literalWithSingleQuotedStringParameter.py @@ -0,0 +1,8 @@ +from typing import Literal + + +def f(x: Literal['xxx']) -> None: + pass + + +f() diff --git a/python/testData/completion/literalType/nestedArgumentLists.py b/python/testData/completion/literalType/nestedArgumentLists.py new file mode 100644 index 000000000000..845dfbe3e46b --- /dev/null +++ b/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("")) \ No newline at end of file diff --git a/python/testData/completion/literalType/varargs.py b/python/testData/completion/literalType/varargs.py new file mode 100644 index 000000000000..085af019b273 --- /dev/null +++ b/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("") diff --git a/python/testSrc/com/jetbrains/python/PyLiteralTypeCompletionTest.kt b/python/testSrc/com/jetbrains/python/PyLiteralTypeCompletionTest.kt new file mode 100644 index 000000000000..42875d175dcf --- /dev/null +++ b/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) + } +} \ No newline at end of file