diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt index 152c08ea5562..c1fd71a2d9ff 100644 --- a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt @@ -13,6 +13,7 @@ import com.intellij.openapi.project.Project import com.intellij.patterns.PlatformPatterns import com.intellij.util.ProcessingContext import com.intellij.xdebugger.impl.ui.tree.nodes.XValueNodeImpl +import com.jetbrains.python.PyTokenTypes import com.jetbrains.python.debugger.PyDebugValue import com.jetbrains.python.debugger.state.PyRuntime import com.jetbrains.python.debugger.values.DataFrameDebugValue @@ -34,9 +35,10 @@ private fun postProcessingChildren(completionResultData: CompletionResultData, return when (completionResultData.completionType) { PyRuntimeCompletionType.DATA_FRAME_COLUMNS -> { val project = parameters.editor.project ?: return emptyList() - processDataFrameColumns(candidate.psiName, + val needValidatorCheck = (candidate.pyQualifiedExpressionList.lastOrNull()?.delimiter ?: candidate.psiName.delimiter) != PyTokenTypes.LBRACKET + processDataFrameColumns(candidate.psiName.pyQualifiedName, completionResultData.setOfCompletionItems, - candidate.needValidatorCheck, + needValidatorCheck, parameters.position, project, true) @@ -52,12 +54,12 @@ interface PyRuntimeCompletionRetrievalService { */ fun canComplete(parameters: CompletionParameters): Boolean - fun extractItemsForCompletion(result: Pair>?, + fun extractItemsForCompletion(result: Pair>?, candidate: PyObjectCandidate): CompletionResultData? { val (node, listOfCalls) = result ?: return null val debugValue = node.valueContainer if (debugValue is DataFrameDebugValue) { - val dfColumns = completePandasDataFrameColumns(debugValue.treeColumns, listOfCalls) ?: return null + val dfColumns = completePandasDataFrameColumns(debugValue.treeColumns, listOfCalls.map { it.pyQualifiedName }) ?: return null return CompletionResultData(dfColumns, PyRuntimeCompletionType.DATA_FRAME_COLUMNS) } computeChildrenIfNeeded(node) @@ -119,11 +121,10 @@ abstract class AbstractRuntimeCompletionContributor : CompletionContributor(), D abstract fun getCompletionRetrievalService(project: Project): PyRuntimeCompletionRetrievalService } -private fun createCompletionResultSet( - retrievalService: PyRuntimeCompletionRetrievalService, - runtimeService: PyRuntime, - parameters: CompletionParameters, -): List { + +fun createCompletionResultSet(retrievalService: PyRuntimeCompletionRetrievalService, + runtimeService: PyRuntime, + parameters: CompletionParameters): List { if (!retrievalService.canComplete(parameters)) return emptyList() val project = parameters.editor.project ?: return emptyList() val treeNodeList = runtimeService.getGlobalPythonVariables(parameters.originalFile.virtualFile, project, parameters.editor) @@ -132,15 +133,17 @@ private fun createCompletionResultSet( return ApplicationUtil.runWithCheckCanceled(Callable { return@Callable pyObjectCandidates.flatMap { candidate -> - val parentNode = getParentNodeByName(treeNodeList, candidate.psiName) + val parentNode = getParentNodeByName(treeNodeList, candidate.psiName.pyQualifiedName) val valueContainer = parentNode?.valueContainer - /** - * Don't need to send requests to jupyter server about Python's module, - * because LegacyCompletionContributor provide completion items for Python's modules - * @see com.intellij.codeInsight.completion.LegacyCompletionContributor - */ - if (valueContainer is PyDebugValue && valueContainer.type == "module") return@flatMap emptyList() - + if (valueContainer is PyDebugValue) { + /** + * Don't need to send requests to jupyter server about Python's module, + * because LegacyCompletionContributor provide completion items for Python's modules + * @see com.intellij.codeInsight.completion.LegacyCompletionContributor + */ + if (valueContainer.type == "module") return@flatMap emptyList() + if (checkDelimiterByType(valueContainer.qualifiedType, candidate.psiName.delimiter)) return@flatMap emptyList() + } getSetOfChildrenByListOfCall(parentNode, candidate.pyQualifiedExpressionList) .let { retrievalService.extractItemsForCompletion(it, candidate) } ?.let { postProcessingChildren(it, candidate, parameters) } diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt index 8a9f70dae2bc..319c7bcde870 100644 --- a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt @@ -14,6 +14,7 @@ import com.intellij.openapi.project.Project import com.intellij.openapi.util.TextRange import com.intellij.openapi.util.text.StringUtil import com.intellij.psi.PsiElement +import com.intellij.psi.tree.IElementType import com.intellij.psi.util.PsiTreeUtil import com.intellij.psi.util.elementType import com.intellij.xdebugger.impl.ui.tree.XDebuggerTreeListener @@ -24,6 +25,7 @@ import com.intellij.xdebugger.impl.ui.tree.nodes.XValueNodeImpl import com.jetbrains.python.PyBundle import com.jetbrains.python.PyTokenTypes import com.jetbrains.python.PythonLanguage +import com.jetbrains.python.debugger.PyDebugValue import com.jetbrains.python.debugger.PyXValueGroup import com.jetbrains.python.debugger.pydev.ProcessDebugger import com.jetbrains.python.debugger.values.DataFrameDebugValue @@ -33,16 +35,21 @@ import com.jetbrains.python.psi.impl.PyStringLiteralExpressionImpl import java.util.concurrent.CompletableFuture import javax.swing.tree.TreeNode +data class PyQualifiedExpressionItem(val pyQualifiedName: String, val delimiter: IElementType) + /** * This data class stores information about a possible python object. * @param psiName - name of PsiElement, that could be a python object - * @param needValidatorCheck - boolean flag to switch on check for validating completion items name - * @param pyQualifiedExpressionList - represents a qualified expression in the list * LanguageNamesValidation.INSTANCE.forLanguage(PythonLanguage.getInstance()).isIdentifier() + * @param pyQualifiedExpressionList - represents a qualified expression in the list + * @see PyQualifiedExpressionItem * - * An example "a.b.c": psiName = "a", needValidatorCheck = false, pyQualifiedExpressionList = ["b","c"] + * An example "a.b.c.": + * psiName = PyQualifiedExpressionItem("a",PyTokenTypes.DOT), + * pyQualifiedExpressionList = [PyQualifiedExpressionItem("b",PyTokenTypes.DOT),PyQualifiedExpressionItem("c",PyTokenTypes.DOT)] */ -data class PyObjectCandidate(val psiName: String, val needValidatorCheck: Boolean, val pyQualifiedExpressionList: List) +data class PyObjectCandidate(val psiName: PyQualifiedExpressionItem, + val pyQualifiedExpressionList: List) // Temporary priority value to control order in CompletionResultSet (DS-3746) @@ -51,9 +58,8 @@ private const val RUNTIME_COMPLETION_PRIORITY = 100.0 fun getCompleteAttribute(parameters: CompletionParameters): List { val callInnerReferenceExpression = getCallInnerReferenceExpression(parameters) - val (currentElement, needValidatorCheck) = - findCompleteAttribute(parameters) - ?: return getPossibleObjectsDataFrame(parameters, callInnerReferenceExpression) + val (currentElement, lastDelimiter) = findCompleteAttribute(parameters) ?: return getPossibleObjectsDataFrame(parameters, + callInnerReferenceExpression) return when (currentElement) { is PyCallExpression, is PyParenthesizedExpression -> { @@ -61,10 +67,10 @@ fun getCompleteAttribute(parameters: CompletionParameters): List buildList { val parentLambdaExpression = collectParentOfLambdaExpression(currentElement, callInnerReferenceExpression) - parentLambdaExpression?.let { - add(createPandasDataFrameCandidate(it, needValidatorCheck)) + parentLambdaExpression?.let { expression -> + createPyObjectCandidate(expression, lastDelimiter)?.let { add(it) } } - add(createPandasDataFrameCandidate(currentElement, needValidatorCheck)) + createPyObjectCandidate(currentElement, lastDelimiter)?.let { add(it) } } else -> { emptyList() @@ -72,36 +78,56 @@ fun getCompleteAttribute(parameters: CompletionParameters): List() +private fun getPyElementText(child: PsiElement?): String? { + child ?: return null + if (child.elementType != PyTokenTypes.DOT && child.elementType != PyTokenTypes.LBRACKET && child.elementType != PyTokenTypes.RBRACKET) { + return when (child) { + is PyStringLiteralExpression -> { + return (child as PyStringLiteralExpressionImpl).stringValue + } + is PyStringElement -> { + child.content + } + else -> { + child.text + } + } + } + return null +} + +private fun createPyObjectCandidate(psiElement: PsiElement, lastDelimiter: IElementType): PyObjectCandidate? { + val names = mutableListOf() + val delimiter = mutableListOf() val firstChild: PsiElement = PsiTreeUtil.getDeepestFirst(psiElement) - val lastChild = PsiTreeUtil.getDeepestLast(psiElement) if (firstChild != lastChild) { var child = PsiTreeUtil.nextLeaf(firstChild) while (child != null) { - if (child.elementType != PyTokenTypes.DOT && child.elementType != PyTokenTypes.LBRACKET && child.elementType != PyTokenTypes.RBRACKET) { - when (child) { - is PyStringLiteralExpression -> { - columns.add((child as PyStringLiteralExpressionImpl).stringValue) - } - is PyStringElement -> { - columns.add(child.content) - } - else -> { - columns.add(child.text) - } - } + if (child.elementType == PyTokenTypes.DOT || child.elementType == PyTokenTypes.LBRACKET) { + child.elementType?.let { delimiter.add(it) } + child = PsiTreeUtil.nextLeaf(child) + continue } + getPyElementText(child)?.let { names.add(it) } if (child == lastChild) { break } child = PsiTreeUtil.nextLeaf(child) } - } - return PyObjectCandidate(firstChild.text, needValidatorCheck, columns) + if (delimiter.size == names.size) { + val firstDelimiter = delimiter.removeAt(0) + delimiter.add(lastDelimiter) + return PyObjectCandidate(PyQualifiedExpressionItem(firstChild.text, firstDelimiter), + names.zip(delimiter).map { pair -> + PyQualifiedExpressionItem(pair.first, pair.second) + }) + } + return null + } + return PyObjectCandidate(PyQualifiedExpressionItem(firstChild.text, lastDelimiter), emptyList()) } private fun getPossibleObjectsDataFrame(parameters: CompletionParameters, @@ -109,22 +135,25 @@ private fun getPossibleObjectsDataFrame(parameters: CompletionParameters, return setOfNotNull(callInnerReferenceExpression?.text, getSliceSubscriptionReferenceExpression(parameters)?.text, getAttributeReferenceExpression(parameters)?.text).map { - PyObjectCandidate(it, false, emptyList()) + PyObjectCandidate(PyQualifiedExpressionItem(it, PyTokenTypes.LBRACKET), emptyList()) } } -private fun findCompleteAttribute(parameters: CompletionParameters): Pair? { - var needCheck = true +private fun findCompleteAttribute(parameters: CompletionParameters): Pair? { var element = parameters.position + val delimiter: IElementType? if (element.prevSibling?.elementType != PyTokenTypes.DOT && element.parent?.prevSibling?.elementType == PyTokenTypes.LBRACKET) { - needCheck = false + delimiter = element.parent?.prevSibling?.elementType element = element.parent } + else { + delimiter = element.prevSibling?.elementType + } val exactElement = element.prevSibling?.prevSibling return when { - exactElement != null -> Pair(exactElement, needCheck) + exactElement != null && delimiter != null -> Pair(exactElement, delimiter) else -> null } @@ -296,7 +325,21 @@ fun getParentNodeByName(children: List, psiName: String): XValueNodeIm return null } -fun getSetOfChildrenByListOfCall(valueNode: XValueNodeImpl?, listOfCall: List): Pair>? { +val typeToDelimiter = mapOf( + "polars.internals.dataframe.frame.DataFrame" to setOf(PyTokenTypes.LBRACKET), + "polars.dataframe.frame.DataFrame" to setOf(PyTokenTypes.LBRACKET), + "pandas.core.frame.DataFrame" to setOf(PyTokenTypes.LBRACKET, PyTokenTypes.DOT) +) + +fun checkDelimiterByType(qualifiedType: String?, delimiter: IElementType): Boolean { + qualifiedType ?: return false + val delimiters = typeToDelimiter[qualifiedType] + if (delimiters != null && delimiter !in delimiters) return true + return false +} + +fun getSetOfChildrenByListOfCall(valueNode: XValueNodeImpl?, + listOfCall: List): Pair>? { var currentNode = valueNode ?: return null listOfCall.forEachIndexed { index, call -> @@ -306,9 +349,13 @@ fun getSetOfChildrenByListOfCall(valueNode: XValueNodeImpl?, listOfCall: List { computeChildrenIfNeeded(currentNode) - currentNode = extractChildByName(currentNode.children, call) ?: return null + currentNode = extractChildByName(currentNode.children, call.pyQualifiedName) ?: return null } } + val valueContainer = currentNode.valueContainer + if (valueContainer is PyDebugValue) { + if (checkDelimiterByType(valueContainer.qualifiedType, call.delimiter)) return null + } } return Pair(currentNode, emptyList()) }