From e94ddb7777ccef91c2c6420b8351d89287f50d02 Mon Sep 17 00:00:00 2001 From: "Olga.Lavrichenko" Date: Fri, 22 Sep 2023 17:53:24 +0000 Subject: [PATCH] DS-4881 Jupyter runtime completion with global variables names by prefix GitOrigin-RevId: cc782a0ab2d0993df10a81a7782f450dc247e91e --- .../PyRuntimeCompletionRetrievalService.kt | 6 +- .../completion/PyRuntimeCompletionUtils.kt | 64 ++++++++++++++----- 2 files changed, 53 insertions(+), 17 deletions(-) diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt index 73ac1fa99a05..deaab7f34c20 100644 --- a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt @@ -111,7 +111,7 @@ private fun postProcessingChildren(completionResultData: CompletionResultData, } private fun proceedPyValueChildrenNames(childrenNodes: Set, - stringPresentation: String, + stringPresentation: String?, ignoreML: Boolean = true): List { return childrenNodes.map { val lookupElement = LookupElementBuilder.create(it).withTypeText(stringPresentation).withIcon( @@ -247,6 +247,10 @@ fun createCompletionResultSet(retrievalService: PyRuntimeCompletionRetrievalServ return ApplicationUtil.runWithCheckCanceled(Callable { return@Callable pyObjectCandidates.flatMap { candidate -> + if (candidate.psiName.delimiter == null) { + return@flatMap getNodesByPrefix(treeNodeList, candidate.psiName.pyQualifiedName, + parameters.completionType).flatMap { proceedPyValueChildrenNames(setOf(it), null) } + } val parentNode = getParentNodeByName(treeNodeList, candidate.psiName.pyQualifiedName, parameters.completionType) val valueContainer = parentNode?.valueContainer if (valueContainer is PyDebugValue) { diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt index bb9e4fc82723..6523c2b34831 100644 --- a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt @@ -11,6 +11,7 @@ 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.refactoring.suggested.startOffset import com.intellij.xdebugger.impl.ui.tree.XDebuggerTreeListener import com.intellij.xdebugger.impl.ui.tree.nodes.XDebuggerTreeNode import com.intellij.xdebugger.impl.ui.tree.nodes.XValueContainerNode @@ -27,7 +28,7 @@ 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) +data class PyQualifiedExpressionItem(val pyQualifiedName: String, val delimiter: IElementType?) /** * This data class stores information about a possible python object. @@ -158,10 +159,20 @@ private fun getPossibleObjectsDataFrame(parameters: CompletionParameters, parseMethodsWithArguments(callExpression)?.let { return it } - return setOfNotNull(callInnerReferenceExpression?.text, - getSliceSubscriptionReferenceExpression(parameters)?.text, - getAttributeReferenceExpression(parameters)?.text).map { - PyObjectCandidate(PyQualifiedExpressionItem(it, PyTokenTypes.LBRACKET), emptyList()) + + return setOfNotNull( + parameters.position, + callInnerReferenceExpression, + getSliceSubscriptionReferenceExpression(parameters), + getAttributeReferenceExpression(parameters) + ).map { + if (it.elementType == PyTokenTypes.IDENTIFIER) { + PyObjectCandidate( + PyQualifiedExpressionItem(it.text.substring(0, parameters.offset - it.startOffset), null), emptyList()) + } + else { + PyObjectCandidate(PyQualifiedExpressionItem(it.text, PyTokenTypes.LBRACKET), emptyList()) + } } } @@ -195,9 +206,9 @@ private data class RuntimeCompletionMethods(val requiredTypes: List?, va private val moduleToMethods = mapOf( "polars" to RuntimeCompletionMethods(listOf("polars.dataframe.frame.DataFrame", "polars.internals.dataframe.frame.DataFrame"), - listOf("any", "approx_unique", "avg", "arg_sort_by", "by_name", "col", "count", "cumsum", "exclude", "first", - "from_epoch", "groups", "head", "implode", "last", "mean", "median", "min", "max", "n_unique", - "quantile", "std", "tail", "sum")), + listOf("any", "approx_unique", "avg", "arg_sort_by", "by_name", "col", "count", "cumsum", "exclude", + "first", "from_epoch", "groups", "head", "implode", "last", "mean", "median", "min", "max", + "n_unique", "quantile", "std", "tail", "sum")), ) fun parseMethodsWithArguments(callExpression: PyCallExpression?): List? { @@ -343,14 +354,12 @@ internal fun getParentNodeByName(children: List, psiName: String, comp * Firstly, looking through loaded variables and if not found - load values inside the group (make a request to jupyter server). */ val globalVariables = children.filterIsInstance() - globalVariables.forEach { node -> - if (node.name == psiName) { - return node - } - } + globalVariables.firstOrNull { it.name == psiName }?.let { return it } + if (completionType == CompletionType.BASIC) return null - val specialVariables = children.filterIsInstance() - .filter { node -> (node.valueContainer as PyXValueGroup).groupType == ProcessDebugger.GROUP_TYPE.SPECIAL } + val specialVariables = children.filterIsInstance().filter { node -> + (node.valueContainer as PyXValueGroup).groupType == ProcessDebugger.GROUP_TYPE.SPECIAL + } specialVariables.forEach { node -> computeChildrenIfNeeded(node) extractChildByName(node, node.loadedChildren, psiName)?.let { @@ -360,6 +369,28 @@ internal fun getParentNodeByName(children: List, psiName: String, comp return null } +private fun prefixMatch(node: TreeNode, result: MutableList, prefix: String) { + (node as? XValueNodeImpl)?.name?.let { + if (it.startsWith(prefix)) { + result.add(it) + } + } +} + +internal fun getNodesByPrefix(children: List, prefix: String, completionType: CompletionType): List { + val result = mutableListOf() + children.forEach { prefixMatch(it, result, prefix) } + + if (completionType == CompletionType.BASIC) return result + children.forEach { node -> + if (node is XValueGroupNodeImpl && (node.valueContainer as PyXValueGroup).groupType == ProcessDebugger.GROUP_TYPE.SPECIAL) { + computeChildrenIfNeeded(node) + node.children.forEach { prefixMatch(it, result, prefix) } + } + } + return result +} + internal val typeToDelimiter = mapOf( "polars.internals.dataframe.frame.DataFrame" to setOf(PyTokenTypes.LBRACKET), "polars.dataframe.frame.DataFrame" to setOf(PyTokenTypes.LBRACKET), @@ -367,7 +398,8 @@ internal val typeToDelimiter = mapOf( "builtins.dict" to setOf(PyTokenTypes.LBRACKET) ) -internal fun checkDelimiterByType(qualifiedType: String?, delimiter: IElementType): Boolean { +internal fun checkDelimiterByType(qualifiedType: String?, delimiter: IElementType?): Boolean { + delimiter ?: return false qualifiedType ?: return false val delimiters = typeToDelimiter[qualifiedType] return delimiters != null && delimiter !in delimiters