diff --git a/python/pluginResources/intellij.python.community.impl.xml b/python/pluginResources/intellij.python.community.impl.xml index 655174bad64d..b932f62f9dfe 100644 --- a/python/pluginResources/intellij.python.community.impl.xml +++ b/python/pluginResources/intellij.python.community.impl.xml @@ -142,9 +142,9 @@ - + diff --git a/python/src/com/jetbrains/python/codeInsight/completion/AbstractRuntimeCompletionContributor.kt b/python/src/com/jetbrains/python/codeInsight/completion/AbstractRuntimeCompletionContributor.kt new file mode 100644 index 000000000000..2d9db3bebbd3 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/AbstractRuntimeCompletionContributor.kt @@ -0,0 +1,86 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion + +import com.intellij.codeInsight.completion.* +import com.intellij.openapi.components.service +import com.intellij.openapi.progress.ProgressManager +import com.intellij.openapi.project.DumbAware +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.TextRange +import com.intellij.patterns.PlatformPatterns +import com.intellij.util.ProcessingContext +import com.jetbrains.python.codeInsight.completion.runtime.patProvider.RemoteFilePathRetrievalService +import com.jetbrains.python.debugger.state.PyRuntime +import com.jetbrains.python.psi.PyStringElement + +abstract class AbstractRuntimeCompletionContributor : CompletionContributor(), DumbAware { + override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) { + val project = parameters.editor.project ?: return + if (parameters.completionType == CompletionType.CLASS_NAME) return + + val context = ProcessingContext() + if (!PlatformPatterns.psiElement().accepts(parameters.position, context)) return + + ProgressManager.checkCanceled() + + val service: PyRuntimeCompletionRetrievalService = getCompletionRetrievalService(project) + if (!service.canComplete(parameters)) return + + fillCompletionVariantsFromRuntime(project, service, parameters, result) + } + + abstract fun getRuntimeEnvService(project: Project): PyRuntime + + abstract fun getCompletionRetrievalService(project: Project): PyRuntimeCompletionRetrievalService + + private fun fillCompletionVariantsFromRuntime( + project: Project, + service: PyRuntimeCompletionRetrievalService, + parameters: CompletionParameters, + result: CompletionResultSet, + ) { + + val runtimeResults: MutableMap = + PyRuntimeCompletionUtils.createCompletionResultSet(service, getRuntimeEnvService(project), parameters) + .associateByTo(hashMapOf(), { it.lookupString }, + { RuntimeLookupElement(it, createCustomMatcher(parameters, result)) }) + + if (runtimeResults.isEmpty()) { + val remoteFileResults = project.service().retrieveRemoteFileLookupElements(parameters) + runtimeResults.putAll(remoteFileResults) + } + + // In general, [createCompletionResultSet] returns an empty list in two cases: + // * If there is no runtime. In that case, it's better to return early to not waste CPU on runRemainingContributors and + // hash table access, even though these operations are fast. + // * If there is nothing found. In that case, it's better to return early again, because there is nothing to add to the result. + // * Other very improbable cases like the absence of the project assigned to the editor, which are handled in a defensive manner. + if (runtimeResults.isEmpty()) return + + if (!result.isStopped) { + result.runRemainingContributors(parameters) { item -> + if (runtimeResults.remove(item.lookupElement.lookupString) != null) { + val prioritizedCompletionResult = item.withLookupElement(PyRuntimeCompletionUtils.createPrioritizedLookupElement(item.lookupElement, true)) + result.withPrefixMatcher(item.prefixMatcher).passResult(prioritizedCompletionResult) + } + else { + result.passResult(item) + } + } + } + + runtimeResults.values.forEach { + result.withPrefixMatcher(it.prefix).addElement(it.lookupElement) + } + } + + private fun createCustomMatcher(parameters: CompletionParameters, result: CompletionResultSet): PrefixMatcher { + val currentElement = parameters.position + if (currentElement is PyStringElement) { + val newPrefix = TextRange.create(currentElement.contentRange.startOffset, + parameters.offset - currentElement.textRange.startOffset).substring(currentElement.text) + return PlainPrefixMatcher(newPrefix) + } + return PlainPrefixMatcher(result.prefixMatcher.prefix) + } +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/CompletionResultData.kt b/python/src/com/jetbrains/python/codeInsight/completion/CompletionResultData.kt new file mode 100644 index 000000000000..92a53e9d0cd2 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/CompletionResultData.kt @@ -0,0 +1,11 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion + +/** + * @param completionType - type of completion items to choose post-processing function in the future + */ +data class CompletionResultData( + val setOfCompletionItems: Set, + val completionType: PyRuntimeCompletionType, + val referenceString: String, +) \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyObjectCandidate.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyObjectCandidate.kt new file mode 100644 index 000000000000..76c8997cff8e --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyObjectCandidate.kt @@ -0,0 +1,20 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion + +/** + * This data class stores information about a possible python object. + * @param psiName - name of PsiElement, that could be a python object + * LanguageNamesValidation.INSTANCE.forLanguage(PythonLanguage.getInstance()).isIdentifier() + * @param pyQualifiedExpressionList - represents a qualified expression in the list + * @param requiredTypes - list of the required objects type for additional check + * @see PyQualifiedExpressionItem + * + * An example "a.b.c.": + * psiName = PyQualifiedExpressionItem("a",PyTokenTypes.DOT), + * pyQualifiedExpressionList = [PyQualifiedExpressionItem("b",PyTokenTypes.DOT),PyQualifiedExpressionItem("c",PyTokenTypes.DOT)] + */ +data class PyObjectCandidate( + val psiName: PyQualifiedExpressionItem, + val pyQualifiedExpressionList: List, + val requiredTypes: List? = null, +) \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyQualifiedExpressionItem.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyQualifiedExpressionItem.kt new file mode 100644 index 000000000000..2c3dc2612a7b --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyQualifiedExpressionItem.kt @@ -0,0 +1,6 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion + +import com.intellij.psi.tree.IElementType + +data class PyQualifiedExpressionItem(val pyQualifiedName: String, val delimiter: IElementType?) \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt index 58aac7174b3e..b63ef4de5a62 100644 --- a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionRetrievalService.kt @@ -1,150 +1,13 @@ // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.jetbrains.python.codeInsight.completion -import com.intellij.codeInsight.completion.* -import com.intellij.codeInsight.lookup.LookupElement -import com.intellij.codeInsight.lookup.LookupElementBuilder -import com.intellij.lang.LanguageNamesValidation -import com.intellij.openapi.application.ex.ApplicationUtil -import com.intellij.openapi.components.service -import com.intellij.openapi.progress.ProgressManager -import com.intellij.openapi.project.DumbAware -import com.intellij.openapi.project.Project -import com.intellij.openapi.util.text.StringUtil -import com.intellij.patterns.PlatformPatterns -import com.intellij.psi.PsiElement -import com.intellij.ui.IconManager -import com.intellij.ui.PlatformIcons -import com.intellij.util.ProcessingContext +import com.intellij.codeInsight.completion.CompletionParameters +import com.intellij.codeInsight.completion.CompletionType 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.state.PyRuntime import com.jetbrains.python.debugger.values.DataFrameDebugValue import com.jetbrains.python.debugger.values.completeDataFrameColumns -import com.jetbrains.python.psi.PyStringElement -import java.util.concurrent.Callable -enum class PyRuntimeCompletionType { - DYNAMIC_CLASS, DATA_FRAME_COLUMNS, DICT_KEYS -} - - -/** - * @param completionType - type of completion items to choose post-processing function in the future - */ -data class CompletionResultData(val setOfCompletionItems: Set, - val completionType: PyRuntimeCompletionType, - val referenceString: String) - -private fun processDataFrameColumns(columns: Set, - needValidatorCheck: Boolean, - elementOnPosition: PsiElement, - project: Project, - stringPresentation: String, - ignoreML: Boolean = true): List { - val validator = LanguageNamesValidation.INSTANCE.forLanguage(PythonLanguage.getInstance()) - return columns.mapNotNull { column -> - when { - !needValidatorCheck && elementOnPosition !is PyStringElement -> { - "'${StringUtil.escapeStringCharacters(column.length, column, "'", false, StringBuilder())}'" - } - !needValidatorCheck -> StringUtil.escapeStringCharacters(column.length, column, "'\"", false, StringBuilder()) - validator.isIdentifier(column, project) -> column - else -> null - }?.let { - val lookupElement = LookupElementBuilder.create(it).withTypeText(stringPresentation).withIcon( - IconManager.getInstance().getPlatformIcon(PlatformIcons.Parameter)) - createPrioritizedLookupElement(lookupElement, ignoreML) - } - } -} - -/** - * @param lookupElement - completion item - * @param prefix - prefix for prefix matcher corresponding to lookupElement - */ -data class RuntimeLookupElement(val lookupElement: LookupElement, val prefix: PrefixMatcher) - -interface RemoteFilePathRetrievalService { - /** - * This function returns a map where the key is file name and value - RuntimeLookupElement. - * @see RuntimeLookupElement - */ - fun retrieveRemoteFileLookupElements(parameters: CompletionParameters): Map -} - -class DummyRemoteFilePathRetrievalService : RemoteFilePathRetrievalService { - override fun retrieveRemoteFileLookupElements(parameters: CompletionParameters): Map = emptyMap() -} - -private fun postProcessingChildren(completionResultData: CompletionResultData, - candidate: PyObjectCandidate, - parameters: CompletionParameters): List { - return when (completionResultData.completionType) { - PyRuntimeCompletionType.DATA_FRAME_COLUMNS -> { - val project = parameters.editor.project ?: return emptyList() - val needValidatorCheck = (candidate.pyQualifiedExpressionList.lastOrNull()?.delimiter - ?: candidate.psiName.delimiter) != PyTokenTypes.LBRACKET - processDataFrameColumns(completionResultData.setOfCompletionItems, - needValidatorCheck, - parameters.position, - project, - PyBundle.message("pandas.completion.type.text", completionResultData.referenceString)) - } - PyRuntimeCompletionType.DICT_KEYS -> { - val setOfCompletionItems = completionResultData.setOfCompletionItems.filter { it != "__len__" }.map { - it.removeSurrounding("\'") - }.toSet() - val project = parameters.editor.project ?: return emptyList() - processDataFrameColumns(setOfCompletionItems, - false, - parameters.position, - project, - PyBundle.message("dict.completion.type.text")) - } - PyRuntimeCompletionType.DYNAMIC_CLASS -> proceedPyValueChildrenNames(completionResultData.setOfCompletionItems, - candidate.psiName.pyQualifiedName) - } -} - -private fun proceedPyValueChildrenNames(childrenNodes: Set, - stringPresentation: String?, - ignoreML: Boolean = true): List { - return childrenNodes.map { - val lookupElement = LookupElementBuilder.create(it).withTypeText(stringPresentation).withIcon( - IconManager.getInstance().getPlatformIcon(PlatformIcons.Parameter)) - createPrioritizedLookupElement(lookupElement, ignoreML) - } -} - -/** - * This function returns string presentation of reference expression. - * - * An example: - * ``` - * class B: - * d = { - * "key1" : df - * } - * ``` - * For completion inside `B.d['key1'].` that corresponding to `df` PyDebugValue returns "B.d['key1']" - */ -private fun getReferenceExpression(pyDebugValue: PyDebugValue, nodeName: String?): String { - nodeName ?: return "" - val parent = pyDebugValue.parent ?: return pyDebugValue.name - var referenceName: String = nodeName - - for (parentValue in generateSequence(parent, PyDebugValue::getParent)) { - if (parentValue.qualifiedType in typeToDelimiter.keys) { - referenceName = "${parentValue.name}[$referenceName]" - } - referenceName = "${parentValue.name}.$referenceName" - } - return referenceName -} interface PyRuntimeCompletionRetrievalService { /** @@ -153,8 +16,10 @@ interface PyRuntimeCompletionRetrievalService { */ fun canComplete(parameters: CompletionParameters): Boolean - fun extractItemsForCompletion(result: Pair>?, - candidate: PyObjectCandidate, completionType: CompletionType): CompletionResultData? { + fun extractItemsForCompletion( + result: Pair>?, + candidate: PyObjectCandidate, completionType: CompletionType, + ): CompletionResultData? { val (node, listOfCalls) = result ?: return null val debugValue = node.valueContainer if (debugValue is DataFrameDebugValue) { @@ -162,7 +27,7 @@ interface PyRuntimeCompletionRetrievalService { return CompletionResultData(dfColumns, PyRuntimeCompletionType.DATA_FRAME_COLUMNS, getReferenceExpression(debugValue, node.name)) } if (completionType == CompletionType.BASIC) return null - computeChildrenIfNeeded(node) + PyRuntimeCompletionUtils.computeChildrenIfNeeded(node) if ((debugValue as PyDebugValue).qualifiedType == "builtins.dict") { return CompletionResultData(node.loadedChildren.mapNotNull { (it as? XValueNodeImpl)?.name }.toSet(), PyRuntimeCompletionType.DICT_KEYS, getReferenceExpression(debugValue, node.name)) @@ -170,102 +35,30 @@ interface PyRuntimeCompletionRetrievalService { return CompletionResultData(node.loadedChildren.mapNotNull { (it as? XValueNodeImpl)?.name }.toSet(), PyRuntimeCompletionType.DYNAMIC_CLASS, getReferenceExpression(debugValue, node.name)) } -} + /** + * This function returns string presentation of reference expression. + * + * An example: + * ``` + * class B: + * d = { + * "key1" : df + * } + * ``` + * For completion inside `B.d['key1'].` that corresponding to `df` PyDebugValue returns "B.d['key1']" + */ + private fun getReferenceExpression(pyDebugValue: PyDebugValue, nodeName: String?): String { + nodeName ?: return "" + val parent = pyDebugValue.parent ?: return pyDebugValue.name + var referenceName: String = nodeName -abstract class AbstractRuntimeCompletionContributor : CompletionContributor(), DumbAware { - override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) { - val project = parameters.editor.project ?: return - if (parameters.completionType == CompletionType.CLASS_NAME) return - - val context = ProcessingContext() - if (!PlatformPatterns.psiElement().accepts(parameters.position, context)) return - - ProgressManager.checkCanceled() - - val service: PyRuntimeCompletionRetrievalService = getCompletionRetrievalService(project) - if (!service.canComplete(parameters)) return - - fillCompletionVariantsFromRuntime(project, service, parameters, result) + for (parentValue in generateSequence(parent, PyDebugValue::getParent)) { + if (parentValue.qualifiedType in PyRuntimeCompletionUtils.typeToDelimiter.keys) { + referenceName = "${parentValue.name}[$referenceName]" + } + referenceName = "${parentValue.name}.$referenceName" + } + return referenceName } - - private fun fillCompletionVariantsFromRuntime( - project: Project, - service: PyRuntimeCompletionRetrievalService, - parameters: CompletionParameters, - result: CompletionResultSet, - ) { - - val runtimeResults: MutableMap = - createCompletionResultSet(service, getRuntimeEnvService(project), parameters) - .associateByTo(hashMapOf(), { it.lookupString }, - { RuntimeLookupElement(it, createCustomMatcher(parameters, result)) }) - - if (runtimeResults.isEmpty()) { - val remoteFileResults = project.service().retrieveRemoteFileLookupElements(parameters) - runtimeResults.putAll(remoteFileResults) - } - - // In general, [createCompletionResultSet] returns an empty list in two cases: - // * If there is no runtime. In that case, it's better to return early to not waste CPU on runRemainingContributors and - // hash table access, even though these operations are fast. - // * If there is nothing found. In that case, it's better to return early again, because there is nothing to add to the result. - // * Other very improbable cases like the absence of the project assigned to the editor, which are handled in a defensive manner. - if (runtimeResults.isEmpty()) return - - if (!result.isStopped) { - result.runRemainingContributors(parameters) { item -> - if (runtimeResults.remove(item.lookupElement.lookupString) != null) { - val prioritizedCompletionResult = item.withLookupElement(createPrioritizedLookupElement(item.lookupElement, true)) - result.withPrefixMatcher(item.prefixMatcher).passResult(prioritizedCompletionResult) - } - else { - result.passResult(item) - } - } - } - - runtimeResults.values.forEach { - result.withPrefixMatcher(it.prefix).addElement(it.lookupElement) - } - } - - abstract fun getRuntimeEnvService(project: Project): PyRuntime - - abstract fun getCompletionRetrievalService(project: Project): PyRuntimeCompletionRetrievalService -} - - -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) - ?: return emptyList() - val pyObjectCandidates = getCompleteAttribute(parameters) - - 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) { - /** - * 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, parameters.completionType) - .let { retrievalService.extractItemsForCompletion(it, candidate, parameters.completionType) } - ?.let { postProcessingChildren(it, candidate, parameters) } - ?: emptyList() - } - }, ProgressManager.getInstance().progressIndicator) } \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionType.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionType.kt new file mode 100644 index 000000000000..50785de31cca --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionType.kt @@ -0,0 +1,6 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion + +enum class PyRuntimeCompletionType { + DYNAMIC_CLASS, DATA_FRAME_COLUMNS, DICT_KEYS +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt index 5dd21cb10792..2ce10e09683f 100644 --- a/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyRuntimeCompletionUtils.kt @@ -1,453 +1,545 @@ // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.jetbrains.python.codeInsight.completion -import com.intellij.codeInsight.completion.* +import com.intellij.codeInsight.completion.CompletionParameters +import com.intellij.codeInsight.completion.CompletionType +import com.intellij.codeInsight.completion.LegacyCompletionContributor +import com.intellij.codeInsight.completion.PrioritizedLookupElement import com.intellij.codeInsight.completion.ml.MLRankingIgnorable import com.intellij.codeInsight.lookup.LookupElement +import com.intellij.codeInsight.lookup.LookupElementBuilder import com.intellij.codeInsight.lookup.LookupElementDecorator +import com.intellij.lang.LanguageNamesValidation +import com.intellij.openapi.application.ex.ApplicationUtil import com.intellij.openapi.application.invokeLater -import com.intellij.openapi.util.TextRange +import com.intellij.openapi.progress.ProgressManager +import com.intellij.openapi.project.Project +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.psi.util.startOffset +import com.intellij.ui.IconManager +import com.intellij.ui.PlatformIcons 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 import com.intellij.xdebugger.impl.ui.tree.nodes.XValueGroupNodeImpl 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.state.PyRuntime import com.jetbrains.python.debugger.values.DataFrameDebugValue import com.jetbrains.python.psi.* import com.jetbrains.python.psi.impl.PyPsiUtils import com.jetbrains.python.psi.impl.PyStringLiteralExpressionImpl +import java.util.concurrent.Callable 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 - * LanguageNamesValidation.INSTANCE.forLanguage(PythonLanguage.getInstance()).isIdentifier() - * @param pyQualifiedExpressionList - represents a qualified expression in the list - * @param requiredTypes - list of the required objects type for additional check - * @see PyQualifiedExpressionItem - * - * An example "a.b.c.": - * psiName = PyQualifiedExpressionItem("a",PyTokenTypes.DOT), - * pyQualifiedExpressionList = [PyQualifiedExpressionItem("b",PyTokenTypes.DOT),PyQualifiedExpressionItem("c",PyTokenTypes.DOT)] - */ -data class PyObjectCandidate(val psiName: PyQualifiedExpressionItem, - val pyQualifiedExpressionList: List, - val requiredTypes: List? = null) +internal object PyRuntimeCompletionUtils { + // Temporary priority value to control order in CompletionResultSet (DS-3746) + private const val RUNTIME_COMPLETION_PRIORITY = 100.0 -// Temporary priority value to control order in CompletionResultSet (DS-3746) -private const val RUNTIME_COMPLETION_PRIORITY = 100.0 + internal fun getCompleteAttribute(parameters: CompletionParameters): List { -internal fun getCompleteAttribute(parameters: CompletionParameters): List { + val callInnerReferenceExpression = getCallInnerReferenceExpression(parameters) + val (currentElement, lastDelimiter) = findCompleteAttribute(parameters) ?: return getPossibleObjectsDataFrame(parameters, + callInnerReferenceExpression) - val callInnerReferenceExpression = getCallInnerReferenceExpression(parameters) - val (currentElement, lastDelimiter) = findCompleteAttribute(parameters) ?: return getPossibleObjectsDataFrame(parameters, - callInnerReferenceExpression) - - return when (currentElement) { - is PyCallExpression, is PyParenthesizedExpression -> { - if (lastDelimiter == PyTokenTypes.DOT) { - val qualifiedElement = PyQualifiedExpressionItem(currentElement.firstChild.text, CALL_DOT) - return listOf(PyObjectCandidate(qualifiedElement, emptyList())) + return when (currentElement) { + is PyCallExpression, is PyParenthesizedExpression -> { + if (lastDelimiter == PyTokenTypes.DOT) { + val qualifiedElement = PyQualifiedExpressionItem(currentElement.firstChild.text, CALL_DOT) + return listOf(PyObjectCandidate(qualifiedElement, emptyList())) + } + emptyList() } - emptyList() - } - is PyExpression -> buildList { - val parentLambdaExpression = collectParentOfLambdaExpression(currentElement, callInnerReferenceExpression) - parentLambdaExpression?.let { expression -> - createPyObjectCandidate(expression, lastDelimiter)?.let { add(it) } - } - createPyObjectCandidate(currentElement, lastDelimiter)?.let { add(it) } - } - else -> { - emptyList() - } - } -} - -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 + is PyExpression -> buildList { + val parentLambdaExpression = collectParentOfLambdaExpression(currentElement, callInnerReferenceExpression) + parentLambdaExpression?.let { expression -> + createPyObjectCandidate(expression, lastDelimiter)?.let { add(it) } + } + createPyObjectCandidate(currentElement, lastDelimiter)?.let { add(it) } } else -> { - child.text + emptyList() } } } - return null -} -private val CALL_DOT = PyElementType("CALL_DOT") -val pattern = listOf(PyTokenTypes.LPAR, PyTokenTypes.RPAR, PyTokenTypes.DOT) - -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) - - val allChildren = generateSequence(firstChild) { child -> - if (child == lastChild) null else PsiTreeUtil.nextLeaf(child) - }.toList() - - var index = 0 - while (index < allChildren.size) { - val elementType = allChildren[index].elementType ?: continue - when (elementType) { - PyTokenTypes.LPAR -> { - if (allChildren.size > index + pattern.size && allChildren.subList(index, index + pattern.size).map { it.elementType } == pattern) { - delimiter.add(CALL_DOT) - index += 3 + 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 } - else { + is PyStringElement -> { + child.content + } + else -> { + child.text + } + } + } + return null + } + + private val CALL_DOT = PyElementType("CALL_DOT") + val pattern = listOf(PyTokenTypes.LPAR, PyTokenTypes.RPAR, PyTokenTypes.DOT) + + 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) + + val allChildren = generateSequence(firstChild) { child -> + if (child == lastChild) null else PsiTreeUtil.nextLeaf(child) + }.toList() + + var index = 0 + while (index < allChildren.size) { + val elementType = allChildren[index].elementType ?: continue + when (elementType) { + PyTokenTypes.LPAR -> { + if (allChildren.size > index + pattern.size && allChildren.subList(index, index + pattern.size).map { it.elementType } == pattern) { + delimiter.add(CALL_DOT) + index += 3 + } + else { + delimiter.add(elementType) + index += 1 + } + } + PyTokenTypes.RBRACKET -> { + if (delimiter.last() != PyTokenTypes.LBRACKET) { + delimiter.add(elementType) + } + index += 1 + } + PyTokenTypes.DOT, PyTokenTypes.LBRACKET -> { delimiter.add(elementType) index += 1 } - } - PyTokenTypes.RBRACKET -> { - if (delimiter.last() != PyTokenTypes.LBRACKET) { - delimiter.add(elementType) + else -> { + val elementText = getPyElementText(allChildren[index]) ?: return null + names.add(elementText) + index += 1 } - index += 1 } - PyTokenTypes.DOT, PyTokenTypes.LBRACKET -> { - delimiter.add(elementType) - index += 1 + } + delimiter.add(lastDelimiter) + + if (names.isNotEmpty() && delimiter.size == names.size) { + val firstDelimiter = delimiter.removeFirst() + val firstName = names.removeFirst() + return PyObjectCandidate(PyQualifiedExpressionItem(firstName, firstDelimiter), names.zip(delimiter).map { pair -> + PyQualifiedExpressionItem(pair.first, pair.second) + }) + } + + return PyObjectCandidate(PyQualifiedExpressionItem(firstChild.text, lastDelimiter), emptyList()) + } + + private fun getPossibleObjectsDataFrame( + parameters: CompletionParameters, + callInnerReferenceExpression: PyExpression?, + ): List { + val callExpression = PsiTreeUtil.getParentOfType(parameters.position, PyCallExpression::class.java) + parseMethodsWithArguments(callExpression)?.let { + return it + } + + 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 -> { - val elementText = getPyElementText(allChildren[index]) ?: return null - names.add(elementText) - index += 1 + else { + PyObjectCandidate(PyQualifiedExpressionItem(it.text, PyTokenTypes.LBRACKET), emptyList()) } } } - delimiter.add(lastDelimiter) - if (names.isNotEmpty() && delimiter.size == names.size) { - val firstDelimiter = delimiter.removeFirst() - val firstName = names.removeFirst() - return PyObjectCandidate(PyQualifiedExpressionItem(firstName, firstDelimiter), names.zip(delimiter).map { pair -> - PyQualifiedExpressionItem(pair.first, pair.second) - }) - } + private fun findCompleteAttribute(parameters: CompletionParameters): Pair? { + var element = parameters.position + val delimiter: IElementType? - return PyObjectCandidate(PyQualifiedExpressionItem(firstChild.text, lastDelimiter), emptyList()) -} - -private fun getPossibleObjectsDataFrame(parameters: CompletionParameters, - callInnerReferenceExpression: PyExpression?): List { - val callExpression = PsiTreeUtil.getParentOfType(parameters.position, PyCallExpression::class.java) - parseMethodsWithArguments(callExpression)?.let { - return it - } - - 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()) + if (element.prevSibling?.elementType != PyTokenTypes.DOT && element.parent?.prevSibling?.elementType == PyTokenTypes.LBRACKET) { + delimiter = element.parent?.prevSibling?.elementType + element = element.parent } else { - PyObjectCandidate(PyQualifiedExpressionItem(it.text, PyTokenTypes.LBRACKET), emptyList()) + delimiter = element.prevSibling?.elementType + } + + val exactElement = element.prevSibling?.prevSibling + return when { + exactElement != null && delimiter != null -> Pair(exactElement, delimiter) + + else -> null } } -} -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) { - delimiter = element.parent?.prevSibling?.elementType - element = element.parent - } - else { - delimiter = element.prevSibling?.elementType + 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) + ?: return emptyList() + val pyObjectCandidates = getCompleteAttribute(parameters) + + 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) { + /** + * Don't need to send requests to jupyter server about Python's module, + * because LegacyCompletionContributor provide completion items for Python's modules + * @see LegacyCompletionContributor + */ + if (valueContainer.type == "module") return@flatMap emptyList() + if (checkDelimiterByType(valueContainer.qualifiedType, candidate.psiName.delimiter)) return@flatMap emptyList() + } + getSetOfChildrenByListOfCall(parentNode, candidate, parameters.completionType) + .let { retrievalService.extractItemsForCompletion(it, candidate, parameters.completionType) } + ?.let { postProcessingChildren(it, candidate, parameters) } + ?: emptyList() + } + }, ProgressManager.getInstance().progressIndicator) } - val exactElement = element.prevSibling?.prevSibling - return when { - exactElement != null && delimiter != null -> Pair(exactElement, delimiter) - - else -> null + private fun postProcessingChildren( + completionResultData: CompletionResultData, + candidate: PyObjectCandidate, + parameters: CompletionParameters, + ): List { + return when (completionResultData.completionType) { + PyRuntimeCompletionType.DATA_FRAME_COLUMNS -> { + val project = parameters.editor.project ?: return emptyList() + val needValidatorCheck = (candidate.pyQualifiedExpressionList.lastOrNull()?.delimiter + ?: candidate.psiName.delimiter) != PyTokenTypes.LBRACKET + processDataFrameColumns(completionResultData.setOfCompletionItems, + needValidatorCheck, + parameters.position, + project, + PyBundle.message("pandas.completion.type.text", completionResultData.referenceString)) + } + PyRuntimeCompletionType.DICT_KEYS -> { + val setOfCompletionItems = completionResultData.setOfCompletionItems.filter { it != "__len__" }.map { + it.removeSurrounding("\'") + }.toSet() + val project = parameters.editor.project ?: return emptyList() + processDataFrameColumns(setOfCompletionItems, + false, + parameters.position, + project, + PyBundle.message("dict.completion.type.text")) + } + PyRuntimeCompletionType.DYNAMIC_CLASS -> proceedPyValueChildrenNames(completionResultData.setOfCompletionItems, + candidate.psiName.pyQualifiedName) + } } -} -// DS-4870 -/** - * This data class retains types and methods about particular expressions associated with certain module usages. - * @see moduleToMethods - */ -private data class RuntimeCompletionMethods(val requiredTypes: List?, val methodNames: List) + private fun processDataFrameColumns( + columns: Set, + needValidatorCheck: Boolean, + elementOnPosition: PsiElement, + project: Project, + stringPresentation: String, + ignoreML: Boolean = true, + ): List { + val validator = LanguageNamesValidation.INSTANCE.forLanguage(PythonLanguage.getInstance()) + return columns.mapNotNull { column -> + when { + !needValidatorCheck && elementOnPosition !is PyStringElement -> { + "'${StringUtil.escapeStringCharacters(column.length, column, "'", false, StringBuilder())}'" + } + !needValidatorCheck -> StringUtil.escapeStringCharacters(column.length, column, "'\"", false, StringBuilder()) + validator.isIdentifier(column, project) -> column + else -> null + }?.let { + val lookupElement = LookupElementBuilder.create(it).withTypeText(stringPresentation).withIcon( + IconManager.getInstance().getPlatformIcon(PlatformIcons.Parameter)) + createPrioritizedLookupElement(lookupElement, ignoreML) + } + } + } + + private fun proceedPyValueChildrenNames( + childrenNodes: Set, + stringPresentation: String?, + ignoreML: Boolean = true, + ): List { + return childrenNodes.map { + val lookupElement = LookupElementBuilder.create(it).withTypeText(stringPresentation).withIcon( + IconManager.getInstance().getPlatformIcon(PlatformIcons.Parameter)) + createPrioritizedLookupElement(lookupElement, ignoreML) + } + } -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")), -) + 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")), + ) -fun parseMethodsWithArguments(callExpression: PyCallExpression?): List? { - val calleeFqn = (callExpression?.callee?.reference?.resolve() as? PyFunction)?.qualifiedName?.split(".") ?: return null - val moduleMethods = moduleToMethods[calleeFqn.firstOrNull()] ?: return null - if (calleeFqn.lastOrNull() in moduleMethods.methodNames) { - var parentExpression = PsiTreeUtil.getParentOfType(callExpression.parent, PyCallExpression::class.java) - val result = mutableListOf() - while (parentExpression != null) { - val psiElement = PsiTreeUtil.getChildOfType(parentExpression, PyReferenceExpression::class.java)?.navigationElement ?: break - val possibleDataFrame = createPyObjectCandidate(psiElement, PyTokenTypes.LPAR) ?: return null - result.add(PyObjectCandidate(PyQualifiedExpressionItem(possibleDataFrame.psiName.pyQualifiedName, PyTokenTypes.LBRACKET), - emptyList(), - moduleMethods.requiredTypes)) - parentExpression = PsiTreeUtil.getParentOfType(parentExpression, PyCallExpression::class.java) + fun parseMethodsWithArguments(callExpression: PyCallExpression?): List? { + val calleeFqn = (callExpression?.callee?.reference?.resolve() as? PyFunction)?.qualifiedName?.split(".") ?: return null + val moduleMethods = moduleToMethods[calleeFqn.firstOrNull()] ?: return null + if (calleeFqn.lastOrNull() in moduleMethods.methodNames) { + var parentExpression = PsiTreeUtil.getParentOfType(callExpression.parent, PyCallExpression::class.java) + val result = mutableListOf() + while (parentExpression != null) { + val psiElement = PsiTreeUtil.getChildOfType(parentExpression, PyReferenceExpression::class.java)?.navigationElement ?: break + val possibleDataFrame = createPyObjectCandidate(psiElement, PyTokenTypes.LPAR) ?: return null + result.add(PyObjectCandidate(PyQualifiedExpressionItem(possibleDataFrame.psiName.pyQualifiedName, PyTokenTypes.LBRACKET), + emptyList(), + moduleMethods.requiredTypes)) + parentExpression = PsiTreeUtil.getParentOfType(parentExpression, PyCallExpression::class.java) + } + return result + } + return null + } + + internal fun getCallInnerReferenceExpression(parameters: CompletionParameters): PyExpression? { + var callExpression: PyCallExpression? = PsiTreeUtil.getParentOfType(parameters.position, PyCallExpression::class.java) ?: return null + var referenceExpression: PyReferenceExpression? = null + + while (true) { + val tmp = PsiTreeUtil.findChildOfType(callExpression, PyReferenceExpression::class.java) + if (tmp == null) break + referenceExpression = tmp + callExpression = PsiTreeUtil.getChildOfType(referenceExpression, PyCallExpression::class.java) + } + if (referenceExpression != null) { + return PyPsiUtils.getFirstQualifier(referenceExpression) + } + return null + } + + private fun getSliceSubscriptionReferenceExpression(parameters: CompletionParameters): PyExpression? { + val expression = PsiTreeUtil.getParentOfType(parameters.position, PySubscriptionExpression::class.java) + val result = PsiTreeUtil.getChildOfType(expression, PyReferenceExpression::class.java) + if (result != null) { + return PyPsiUtils.getFirstQualifier(result) + } + if (expression != null) { + return getSliceSubscriptionReferenceExpressionWithMultiIndex(expression) + } + return null + } + + private fun getSliceSubscriptionReferenceExpressionWithMultiIndex(expression: PyExpression): PyExpression? { + var count = 0 + var child = expression.copy() + + while (PsiTreeUtil.getChildOfType(child, PyReferenceExpression::class.java) == null && child != null) { + count++ + child = PsiTreeUtil.getChildOfType(child, PySubscriptionExpression::class.java) + } + + val result = PsiTreeUtil.getChildOfType(child, PyReferenceExpression::class.java) + if (result != null) { + return PyPsiUtils.getFirstQualifier(result) + } + return null + } + + private fun getAttributeReferenceExpression(parameters: CompletionParameters): PyExpression? { + val result = PsiTreeUtil.getParentOfType(parameters.position, PyReferenceExpression::class.java) + + if (result != null) { + val child = PsiTreeUtil.getChildrenOfType(result, PyExpression::class.java) ?: emptyArray() + if (!child.isEmpty()) { + return child[0] + } + return PyPsiUtils.getFirstQualifier(result) + } + + return null + } + + private fun collectParentOfLambdaExpression(element: PyExpression, callInnerReferenceExpression: PyExpression?): PyExpression? { + if (callInnerReferenceExpression != null) { + val lambdaExpression = PsiTreeUtil.getParentOfType(element, PyLambdaExpression::class.java) + + if (lambdaExpression != null && callInnerReferenceExpression is PyQualifiedExpression) { + if (lambdaExpression.parameterList.parameters.any { it.name == element.name }) { + val callInnerReferencePyExpression = callInnerReferenceExpression as? PyQualifiedExpression + return callInnerReferencePyExpression?.let { PyPsiUtils.getFirstQualifier(it) } + } + } + } + return null + } + + internal fun createPrioritizedLookupElement(lookupElement: LookupElement, ignoreML: Boolean): LookupElement { + val prioritizedElement = PrioritizedLookupElement.withPriority(lookupElement, RUNTIME_COMPLETION_PRIORITY) + if (ignoreML) { + return prioritizedElement.asMLIgnorable() + } + return prioritizedElement + } + + private fun LookupElement.asMLIgnorable(): LookupElement { + return object : LookupElementDecorator(this), MLRankingIgnorable {} + } + + + internal fun computeChildrenIfNeeded(valueNode: XValueContainerNode<*>) { + if (!valueNode.isLeaf && valueNode.loadedChildren.isEmpty()) { + val futureChildrenReady = CompletableFuture() + val listener = object : XDebuggerTreeListener { + override fun childrenLoaded(node: XDebuggerTreeNode, children: MutableList>, last: Boolean) { + if (node == valueNode) { + futureChildrenReady.complete(true) + } + } + } + // Start to listen to tree changes. + valueNode.tree.addTreeListener(listener) + invokeLater { + valueNode.startComputingChildren() + } + // Additional check (children did not empty yet) to prevent race condition. + if (valueNode.loadedChildren.isEmpty()) { + // Wait until children will be added to the tree. + futureChildrenReady.get() + } + // Stop to listen to tree changes. + valueNode.tree.removeTreeListener(listener) + } + } + + private fun extractChildByName(currentNode: XValueContainerNode<*>, childrenNodes: List, name: String): XValueNodeImpl? { + if ((currentNode.valueContainer as? PyDebugValue)?.qualifiedType == "builtins.dict") { + return childrenNodes.firstOrNull { it is XValueNodeImpl && it.name == "'${name}'" } as XValueNodeImpl? + } + return childrenNodes.firstOrNull { it is XValueNodeImpl && it.name == name } as XValueNodeImpl? + } + + internal fun getParentNodeByName(children: List, psiName: String, completionType: CompletionType): XValueNodeImpl? { + /** + * For preventing an extra load of "Special variables". + * 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.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 + } + specialVariables.forEach { node -> + computeChildrenIfNeeded(node) + extractChildByName(node, node.loadedChildren, psiName)?.let { + return it + } + } + 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 } - return null -} -private fun getCallInnerReferenceExpression(parameters: CompletionParameters): PyExpression? { - var callExpression: PyCallExpression? = PsiTreeUtil.getParentOfType(parameters.position, PyCallExpression::class.java) ?: return null - var referenceExpression: PyReferenceExpression? = null + internal 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), + "builtins.dict" to setOf(PyTokenTypes.LBRACKET) + ) - while (true) { - val tmp = PsiTreeUtil.findChildOfType(callExpression, PyReferenceExpression::class.java) - if (tmp == null) break - referenceExpression = tmp - callExpression = PsiTreeUtil.getChildOfType(referenceExpression, PyCallExpression::class.java) - } - if (referenceExpression != null) { - return PyPsiUtils.getFirstQualifier(referenceExpression) - } - return null -} - -private fun getSliceSubscriptionReferenceExpression(parameters: CompletionParameters): PyExpression? { - val expression = PsiTreeUtil.getParentOfType(parameters.position, PySubscriptionExpression::class.java) - val result = PsiTreeUtil.getChildOfType(expression, PyReferenceExpression::class.java) - if (result != null) { - return PyPsiUtils.getFirstQualifier(result) - } - if (expression != null) { - return getSliceSubscriptionReferenceExpressionWithMultiIndex(expression) - } - return null -} - -private fun getSliceSubscriptionReferenceExpressionWithMultiIndex(expression: PyExpression): PyExpression? { - var count = 0 - var child = expression.copy() - - while (PsiTreeUtil.getChildOfType(child, PyReferenceExpression::class.java) == null && child != null) { - count++ - child = PsiTreeUtil.getChildOfType(child, PySubscriptionExpression::class.java) + internal fun checkDelimiterByType(qualifiedType: String?, delimiter: IElementType?): Boolean { + delimiter ?: return false + qualifiedType ?: return false + val delimiters = typeToDelimiter[qualifiedType] + return delimiters != null && delimiter !in delimiters } - val result = PsiTreeUtil.getChildOfType(child, PyReferenceExpression::class.java) - if (result != null) { - return PyPsiUtils.getFirstQualifier(result) - } - return null -} - -private fun getAttributeReferenceExpression(parameters: CompletionParameters): PyExpression? { - val result = PsiTreeUtil.getParentOfType(parameters.position, PyReferenceExpression::class.java) - - if (result != null) { - val child = PsiTreeUtil.getChildrenOfType(result, PyExpression::class.java) ?: emptyArray() - if (!child.isEmpty()) { - return child[0] - } - return PyPsiUtils.getFirstQualifier(result) + internal fun checkRequiredType(qualifiedType: String?, requiredTypes: List?): Boolean { + requiredTypes ?: return true + qualifiedType ?: return false + return requiredTypes.contains(qualifiedType) } - return null -} - -private fun collectParentOfLambdaExpression(element: PyExpression, callInnerReferenceExpression: PyExpression?): PyExpression? { - if (callInnerReferenceExpression != null) { - val lambdaExpression = PsiTreeUtil.getParentOfType(element, PyLambdaExpression::class.java) - - if (lambdaExpression != null && callInnerReferenceExpression is PyQualifiedExpression) { - if (lambdaExpression.parameterList.parameters.any { it.name == element.name }) { - val callInnerReferencePyExpression = callInnerReferenceExpression as? PyQualifiedExpression - return callInnerReferencePyExpression?.let { PyPsiUtils.getFirstQualifier(it) } - } - } - } - return null -} - -internal fun createPrioritizedLookupElement(lookupElement: LookupElement, ignoreML: Boolean): LookupElement { - val prioritizedElement = PrioritizedLookupElement.withPriority(lookupElement, RUNTIME_COMPLETION_PRIORITY) - if (ignoreML) { - return prioritizedElement.asMLIgnorable() - } - return prioritizedElement -} - -private fun LookupElement.asMLIgnorable(): LookupElement { - return object : LookupElementDecorator(this), MLRankingIgnorable {} -} - - -internal fun computeChildrenIfNeeded(valueNode: XValueContainerNode<*>) { - if (!valueNode.isLeaf && valueNode.loadedChildren.isEmpty()) { - val futureChildrenReady = CompletableFuture() - val listener = object : XDebuggerTreeListener { - override fun childrenLoaded(node: XDebuggerTreeNode, children: MutableList>, last: Boolean) { - if (node == valueNode) { - futureChildrenReady.complete(true) + internal fun getSetOfChildrenByListOfCall( + valueNode: XValueNodeImpl?, + candidate: PyObjectCandidate, + completionType: CompletionType, + ): Pair>? { + var currentNode = valueNode ?: return null + val listOfCall = candidate.pyQualifiedExpressionList + listOfCall.forEachIndexed { index, call -> + when (currentNode.valueContainer) { + is DataFrameDebugValue -> { + if (checkRequiredType((currentNode.valueContainer as? DataFrameDebugValue)?.qualifiedType, candidate.requiredTypes)) { + return Pair(currentNode, listOfCall.subList(index, listOfCall.size)) + } + return null + } + else -> { + if (completionType == CompletionType.BASIC) return null + computeChildrenIfNeeded(currentNode) + currentNode = extractChildByName(currentNode, currentNode.children, call.pyQualifiedName) ?: return null } } - } - // Start to listen to tree changes. - valueNode.tree.addTreeListener(listener) - invokeLater { - valueNode.startComputingChildren() - } - // Additional check (children did not empty yet) to prevent race condition. - if (valueNode.loadedChildren.isEmpty()) { - // Wait until children will be added to the tree. - futureChildrenReady.get() - } - // Stop to listen to tree changes. - valueNode.tree.removeTreeListener(listener) - } -} - -private fun extractChildByName(currentNode: XValueContainerNode<*>, childrenNodes: List, name: String): XValueNodeImpl? { - if ((currentNode.valueContainer as? PyDebugValue)?.qualifiedType == "builtins.dict") { - return childrenNodes.firstOrNull { it is XValueNodeImpl && it.name == "'${name}'" } as XValueNodeImpl? - } - return childrenNodes.firstOrNull { it is XValueNodeImpl && it.name == name } as XValueNodeImpl? -} - -internal fun getParentNodeByName(children: List, psiName: String, completionType: CompletionType): XValueNodeImpl? { - /** - * For preventing an extra load of "Special variables". - * 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.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 - } - specialVariables.forEach { node -> - computeChildrenIfNeeded(node) - extractChildByName(node, node.loadedChildren, psiName)?.let { - return it - } - } - 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), - "pandas.core.frame.DataFrame" to setOf(PyTokenTypes.LBRACKET, PyTokenTypes.DOT), - "builtins.dict" to setOf(PyTokenTypes.LBRACKET) -) - -internal fun checkDelimiterByType(qualifiedType: String?, delimiter: IElementType?): Boolean { - delimiter ?: return false - qualifiedType ?: return false - val delimiters = typeToDelimiter[qualifiedType] - return delimiters != null && delimiter !in delimiters -} - -internal fun checkRequiredType(qualifiedType: String?, requiredTypes: List?): Boolean { - requiredTypes ?: return true - qualifiedType ?: return false - return requiredTypes.contains(qualifiedType) -} - -internal fun getSetOfChildrenByListOfCall(valueNode: XValueNodeImpl?, - candidate: PyObjectCandidate, - completionType: CompletionType): Pair>? { - var currentNode = valueNode ?: return null - val listOfCall = candidate.pyQualifiedExpressionList - listOfCall.forEachIndexed { index, call -> - when (currentNode.valueContainer) { - is DataFrameDebugValue -> { - if (checkRequiredType((currentNode.valueContainer as? DataFrameDebugValue)?.qualifiedType, candidate.requiredTypes)) { - return Pair(currentNode, listOfCall.subList(index, listOfCall.size)) - } - return null - } - else -> { - if (completionType == CompletionType.BASIC) return null - computeChildrenIfNeeded(currentNode) - currentNode = extractChildByName(currentNode, currentNode.children, call.pyQualifiedName) ?: return null + val valueContainer = currentNode.valueContainer + if (valueContainer is PyDebugValue) { + if (checkDelimiterByType(valueContainer.qualifiedType, call.delimiter)) return null } } - val valueContainer = currentNode.valueContainer - if (valueContainer is PyDebugValue) { - if (checkDelimiterByType(valueContainer.qualifiedType, call.delimiter)) return null + (currentNode.valueContainer as? PyDebugValue)?.let { + if (checkRequiredType(it.qualifiedType, candidate.requiredTypes)) { + return Pair(currentNode, emptyList()) + } } + return null } - (currentNode.valueContainer as? PyDebugValue)?.let { - if (checkRequiredType(it.qualifiedType, candidate.requiredTypes)) { - return Pair(currentNode, emptyList()) - } - } - return null -} - -internal fun createCustomMatcher(parameters: CompletionParameters, result: CompletionResultSet): PrefixMatcher { - val currentElement = parameters.position - if (currentElement is PyStringElement) { - val newPrefix = TextRange.create(currentElement.contentRange.startOffset, - parameters.offset - currentElement.textRange.startOffset).substring(currentElement.text) - return PlainPrefixMatcher(newPrefix) - } - return PlainPrefixMatcher(result.prefixMatcher.prefix) -} +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/RuntimeCompletionMethods.kt b/python/src/com/jetbrains/python/codeInsight/completion/RuntimeCompletionMethods.kt new file mode 100644 index 000000000000..12594eb82455 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/RuntimeCompletionMethods.kt @@ -0,0 +1,8 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion + +/** + * This data class retains types and methods about particular expressions associated with certain module usages. + * @see moduleToMethods + */ +internal data class RuntimeCompletionMethods(val requiredTypes: List?, val methodNames: List) \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/RuntimeLookupElement.kt b/python/src/com/jetbrains/python/codeInsight/completion/RuntimeLookupElement.kt new file mode 100644 index 000000000000..e9e5c8e5deef --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/RuntimeLookupElement.kt @@ -0,0 +1,11 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion + +import com.intellij.codeInsight.completion.PrefixMatcher +import com.intellij.codeInsight.lookup.LookupElement + +/** + * @param lookupElement - completion item + * @param prefix - prefix for prefix matcher corresponding to lookupElement + */ +data class RuntimeLookupElement(val lookupElement: LookupElement, val prefix: PrefixMatcher) \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/runtime/patProvider/DummyRemoteFilePathRetrievalService.kt b/python/src/com/jetbrains/python/codeInsight/completion/runtime/patProvider/DummyRemoteFilePathRetrievalService.kt new file mode 100644 index 000000000000..e468061a7024 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/runtime/patProvider/DummyRemoteFilePathRetrievalService.kt @@ -0,0 +1,9 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion.runtime.patProvider + +import com.intellij.codeInsight.completion.CompletionParameters +import com.jetbrains.python.codeInsight.completion.RuntimeLookupElement + +class DummyRemoteFilePathRetrievalService : RemoteFilePathRetrievalService { + override fun retrieveRemoteFileLookupElements(parameters: CompletionParameters): Map = emptyMap() +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/completion/runtime/patProvider/RemoteFilePathRetrievalService.kt b/python/src/com/jetbrains/python/codeInsight/completion/runtime/patProvider/RemoteFilePathRetrievalService.kt new file mode 100644 index 000000000000..e9a8afdbc7fc --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/runtime/patProvider/RemoteFilePathRetrievalService.kt @@ -0,0 +1,13 @@ +// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.completion.runtime.patProvider + +import com.intellij.codeInsight.completion.CompletionParameters +import com.jetbrains.python.codeInsight.completion.RuntimeLookupElement + +interface RemoteFilePathRetrievalService { + /** + * This function returns a map where the key is file name and value - RuntimeLookupElement. + * @see com.jetbrains.python.codeInsight.completion.RuntimeLookupElement + */ + fun retrieveRemoteFileLookupElements(parameters: CompletionParameters): Map +} \ No newline at end of file