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()
)
)
}
}
}