PY-56004 Completion for expected literal types in assignments

GitOrigin-RevId: 06cacc62e068d902b174e5208837959cbd92ed88
This commit is contained in:
Petr
2024-04-26 17:29:46 +02:00
committed by intellij-monorepo-bot
parent 39b60e7bee
commit 7965dde4a4
7 changed files with 61 additions and 32 deletions

View File

@@ -4,6 +4,7 @@ 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.PsiElement
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.ui.IconManager
import com.intellij.ui.PlatformIcons
@@ -12,6 +13,7 @@ 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.PyType
import com.jetbrains.python.psi.types.PyTypeUtil
import com.jetbrains.python.psi.types.TypeEvalContext
@@ -22,7 +24,8 @@ class PyLiteralTypeCompletionContributor : CompletionContributor() {
or(
psiElement().withSuperParent(2, PyKeywordArgument::class.java),
psiElement().withSuperParent(2, PyArgumentList::class.java),
psiElement().withSuperParent(2, PySubscriptionExpression::class.java)
psiElement().withSuperParent(2, PySubscriptionExpression::class.java),
psiElement().inside(PyAssignmentStatement::class.java),
),
PyLiteralTypeCompletionProvider()
)
@@ -32,34 +35,51 @@ class PyLiteralTypeCompletionContributor : CompletionContributor() {
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)
val callSiteExpr = PsiTreeUtil.getParentOfType(position, PyCallSiteExpression::class.java)
if (callSiteExpr != null) {
val parent = position.parent
val argumentExpr = if (parent is PyKeywordArgument && parent.valueExpression == position) parent else position
val types = PyCallExpressionHelper
.mapArguments(callSiteExpr, PyResolveContext.defaultContext(typeEvalContext))
.mapNotNull { it.mappedParameters[argumentExpr]?.getArgumentType(typeEvalContext) }
addToResult(position, types, result)
return
}
val assignmentStatement = PsiTreeUtil.skipParentsOfType(position,
PyParenthesizedExpression::class.java,
PyTupleExpression::class.java) as? PyAssignmentStatement
if (assignmentStatement != null) {
val mapping = assignmentStatement.targetsToValuesMapping.find { it.second === position }
if (mapping != null) {
val type = typeEvalContext.getType(mapping.first)
addToResult(position, listOfNotNull(type), result)
}
}
}
private fun addToResult(result: CompletionResultSet, lookupString: String) {
result.addElement(
LookupElementBuilder
.create(lookupString)
.withIcon(IconManager.getInstance().getPlatformIcon(PlatformIcons.Parameter))
)
private fun addToResult(position: PyExpression, possibleTypes: List<PyType>, result: CompletionResultSet) {
val lookupString = if (position is PyStringLiteralExpression)
StringLiteralExpression::getStringValue
else
PsiElement::getText
possibleTypes.asSequence()
.flatMap { PyTypeUtil.toStream(it) }
.filterIsInstance<PyLiteralType>()
.map { it.expression }
.filterIsInstance<PyStringLiteralExpression>()
.forEach {
result.addElement(
PrioritizedLookupElement.withPriority(
LookupElementBuilder
.create(lookupString(it))
.withIcon(IconManager.getInstance().getPlatformIcon(PlatformIcons.Parameter)),
PythonCompletionWeigher.PRIORITY_WEIGHT.toDouble()
)
)
}
}
}

View File

@@ -0,0 +1,5 @@
from typing import Literal
item: Literal[Literal[Literal["1", 2, "3"], "foo"], "5", None]
item=<caret>

View File

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

View File

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

View File

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

View File

@@ -1,7 +1,7 @@
from typing import Literal
def f(x: Literal[22]) -> None:
def f(x: Literal["zzz"]) -> None:
pass

View File

@@ -8,15 +8,19 @@ import com.jetbrains.python.fixtures.PyTestCase
@TestDataPath("\$CONTENT_ROOT/../testData/completion/literalType")
class PyLiteralTypeCompletionTest : PyTestCase() {
fun testInCallExpression() {
doTestCompletionVariantsContains("inCallExpression.py", "1", "2", "3", "\"foo\"", "5", "None")
doTestCompletionVariantsContains("inCallExpression.py", "\"1\"", "\"2\"", "\"foo\"", "\"5\"")
}
fun testInKeywordArgument() {
doTestCompletionVariantsContains("inKeywordArgument.py", "1", "2", "3", "\"foo\"", "5", "None")
doTestCompletionVariantsContains("inKeywordArgument.py", "\"3\"", "\"foo\"", "\"5\"")
}
fun testInSubscriptionExpression() {
doTestCompletionVariantsContains("inSubscriptionExpression.py", "1", "2", "3", "\"foo\"", "5", "None")
doTestCompletionVariantsContains("inSubscriptionExpression.py", "\"1\"", "\"foo\"", "\"5\"")
}
fun testInAssigment() {
doTestCompletionVariantsContains("inAssignment.py", "\"1\"", "\"3\"", "\"foo\"", "\"5\"")
}
fun testInDoubleQuotedString() {
@@ -28,7 +32,7 @@ class PyLiteralTypeCompletionTest : PyTestCase() {
}
fun testLiteralWithSingleParameter() {
doTestCompletionVariantsContains("literalWithSingleParameter.py", "22")
doTestCompletionVariantsContains("literalWithSingleParameter.py", "\"zzz\"")
}
fun testLiteralWithDoubleQuotedStringParameter() {