From 58badfd49287dbf1651c5025a710e9cdceb9f867 Mon Sep 17 00:00:00 2001 From: "Olga.Lavrichenko" Date: Fri, 30 Sep 2022 18:29:21 +0000 Subject: [PATCH] DS-3803 DS-3803 DS-3635 DS-2871 Jupyter autocompletion implicitly executes code in the editor, fixes and moved ConsolePandasColumnNameCompletionContributor to intellij.pycharm.community.iml module GitOrigin-RevId: bae7ec7e06cc96feeb3a8f7ed4eb8d897e289547 --- python/intellij.python.community.impl.iml | 2 + .../messages/PyBundle.properties | 3 +- python/src/META-INF/python-core-common.xml | 209 ++++++++++++------ ...lePandasColumnNameCompletionContributor.kt | 127 +++++++++++ .../completion/PandasCompletionUtils.kt | 195 ++++++++++++++++ 5 files changed, 467 insertions(+), 69 deletions(-) create mode 100644 python/src/com/jetbrains/python/console/completion/ConsolePandasColumnNameCompletionContributor.kt create mode 100644 python/src/com/jetbrains/python/console/completion/PandasCompletionUtils.kt diff --git a/python/intellij.python.community.impl.iml b/python/intellij.python.community.impl.iml index be14917ec435..c0d60927c090 100644 --- a/python/intellij.python.community.impl.iml +++ b/python/intellij.python.community.impl.iml @@ -110,5 +110,7 @@ + + \ No newline at end of file diff --git a/python/pluginResources/messages/PyBundle.properties b/python/pluginResources/messages/PyBundle.properties index 56b1e4d229aa..c3fbbe0f5469 100644 --- a/python/pluginResources/messages/PyBundle.properties +++ b/python/pluginResources/messages/PyBundle.properties @@ -1278,5 +1278,4 @@ python.interpreters.rename.interpreter.dialog.title=Rename Python Interpreter rename.python.interpreter.dialog.provide.name.error.text=Provide a name for Python interpreter. rename.python.interpreter.name.already.exists.error.text=Python interpreter with the same name already exists. - - +pandas.completion.type.text=DataFrame: {0} diff --git a/python/src/META-INF/python-core-common.xml b/python/src/META-INF/python-core-common.xml index 4a3ae8270a94..e2da2f9c8ef8 100644 --- a/python/src/META-INF/python-core-common.xml +++ b/python/src/META-INF/python-core-common.xml @@ -61,7 +61,9 @@ - + @@ -71,7 +73,8 @@ - + @@ -92,6 +95,14 @@ + + + + @@ -111,8 +122,8 @@ - - + + @@ -165,29 +176,49 @@ - + - - - - - - + + + + + + - - - + + + - + @@ -214,12 +245,15 @@ implementationClass="com.jetbrains.python.codeInsight.regexp.PythonVerboseRegexpParserDefinition"/> - + - - + + - + @@ -266,13 +301,17 @@ key="configurable.PySmartKeysOptions.display.name"/> - - - + + + - - + - + @@ -337,8 +378,10 @@ - - + + @@ -364,12 +407,14 @@ - + - - + + @@ -380,22 +425,31 @@ - - - - - - + + + + + + - + - - + + @@ -409,17 +463,24 @@ - - - + + + - - - - + + + + - + - + - - - - + + + + @@ -452,13 +517,17 @@ - + - - - + + + @@ -466,14 +535,16 @@ interface="com.jetbrains.python.console.PyConsoleOptionsProvider" dynamic="true"/> - + - + - + - - + + @@ -674,22 +746,23 @@ - + - + - + - + @@ -699,7 +772,8 @@ - + @@ -723,7 +797,8 @@ - + diff --git a/python/src/com/jetbrains/python/console/completion/ConsolePandasColumnNameCompletionContributor.kt b/python/src/com/jetbrains/python/console/completion/ConsolePandasColumnNameCompletionContributor.kt new file mode 100644 index 000000000000..c80821e17fb0 --- /dev/null +++ b/python/src/com/jetbrains/python/console/completion/ConsolePandasColumnNameCompletionContributor.kt @@ -0,0 +1,127 @@ +package com.jetbrains.python.console.completion + +import com.intellij.codeInsight.completion.* +import com.intellij.execution.ui.RunContentDescriptor +import com.intellij.openapi.application.ex.ApplicationUtil +import com.intellij.openapi.components.service +import com.intellij.openapi.diagnostic.Logger +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.progress.ProgressManager +import com.intellij.openapi.project.DumbAware +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.NlsSafe +import com.intellij.openapi.vfs.VirtualFile +import com.intellij.patterns.PlatformPatterns +import com.intellij.util.ProcessingContext +import com.jetbrains.python.actions.checkIfAvailableAndShowHint +import com.jetbrains.python.actions.getCustomDescriptor +import com.jetbrains.python.actions.getSelectedPythonConsole +import com.jetbrains.python.console.PyExecuteConsoleCustomizer +import com.jetbrains.python.console.PydevConsoleCommunication +import com.jetbrains.python.console.PythonConsoleView +import kotlinx.serialization.SerializationException +import kotlinx.serialization.decodeFromString +import kotlinx.serialization.json.Json +import java.lang.invoke.MethodHandles +import java.util.concurrent.Callable + +class ConsolePandasColumnNameCompletionContributor : CompletionContributor(), DumbAware { + + init { + extend(CompletionType.BASIC, PlatformPatterns.psiElement(), object : CompletionProvider() { + override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) { + val project = parameters.editor.project ?: return + val editor = parameters.editor + val virtualFile = parameters.originalFile.virtualFile + + if (!checkIfAvailableAndShowHint(editor)) return + val existingConsole = getDescriptorIfExist(virtualFile, project, editor) ?: return + val pydevRunner: PythonConsoleView = existingConsole.executionConsole as? PythonConsoleView ?: return + val consoleCommunication = pydevRunner.executeActionHandler?.consoleCommunication as? PydevConsoleCommunication ?: return + if (consoleCommunication.isExecuting) return + + val dataFrameCandidates = getCompleteAttribute(parameters) + ApplicationUtil.runWithCheckCanceled(Callable { + result.addAllElements(dataFrameCandidates.flatMap { candidate -> + val columnsDataFrame = getDataFrameColumns(project, candidate.psiName, consoleCommunication) + processDataFrameColumns(candidate.psiName, columnsDataFrame, candidate.needValidatorCheck, + parameters.position, project, false) + }) + }, ProgressManager.getInstance().progressIndicator) + } + }) + } + + private fun getDescriptorIfExist(virtualFile: VirtualFile?, project: Project, editor: Editor): RunContentDescriptor? { + if (virtualFile != null && PyExecuteConsoleCustomizer.instance.isCustomDescriptorSupported(virtualFile)) { + tryGetCustomDescriptor(project, editor)?.let { return it } + } + return getSelectedPythonConsole(project) + } + + private fun tryGetCustomDescriptor(project: Project, editor: Editor): RunContentDescriptor? { + try { + val (descriptor, _) = getCustomDescriptor(project, editor) + return descriptor + } + catch (e: IllegalStateException) { + return null + } + } + + + private fun getDataFrameColumns(project: Project, name: @NlsSafe String, consoleCommunication: PydevConsoleCommunication): List { + val service = project.service() + return service.getPandasColumns(consoleCommunication, name) + } + + companion object { + val LOG = Logger.getInstance(MethodHandles.lookup().lookupClass()) + const val COMPLETION_LOG_MESSAGE = "Incorrectly created python script with expression: " + } + +} + +interface ConsolePandasColumnNameRetrievalService { + fun getPandasColumns(consoleCommunication: PydevConsoleCommunication, name: @NlsSafe String): List +} + +class DummyConsolePandasColumnNameRetrievalService : ConsolePandasColumnNameRetrievalService { + override fun getPandasColumns(consoleCommunication: PydevConsoleCommunication, name: @NlsSafe String): List = listOf() +} + + +class ConsolePandasColumnNameRetrievalServiceImpl(val project: Project) : ConsolePandasColumnNameRetrievalService { + override fun getPandasColumns(consoleCommunication: PydevConsoleCommunication, name: @NlsSafe String): List { + + val debugValue = consoleCommunication.evaluate(PANDAS_COLUMN_NAMES_CODE.format(name, name), true, true) + return when (debugValue.type) { + "str" -> debugValue.value?.let { parseDebugValue(it) } ?: emptyList() + "NameError" -> emptyList() + "SyntaxError" -> { + ConsolePandasColumnNameCompletionContributor.LOG.error(ConsolePandasColumnNameCompletionContributor.COMPLETION_LOG_MESSAGE + name) + emptyList() + } + else -> emptyList() + } + } + + + private fun parseDebugValue(value: String): List? { + + try { + val map = Json.decodeFromString>>(value) + return map["columns"] + } + catch (_: SerializationException) { + return emptyList() + } + catch (_: IllegalArgumentException) { + return emptyList() + } + } + + companion object { + private const val PANDAS_COLUMN_NAMES_CODE = "__import__('json').dumps({\"columns\":list(%s.columns.get_level_values(0))}) if str(%s.__class__)==\"\" else {}" + } +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/console/completion/PandasCompletionUtils.kt b/python/src/com/jetbrains/python/console/completion/PandasCompletionUtils.kt new file mode 100644 index 000000000000..17dbac059d39 --- /dev/null +++ b/python/src/com/jetbrains/python/console/completion/PandasCompletionUtils.kt @@ -0,0 +1,195 @@ +package com.jetbrains.python.console.completion + +import com.intellij.codeInsight.completion.CompletionParameters +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.project.Project +import com.intellij.openapi.util.text.StringUtil +import com.intellij.psi.PsiElement +import com.intellij.psi.util.PsiTreeUtil +import com.intellij.psi.util.elementType +import com.jetbrains.python.PyBundle +import com.jetbrains.python.PyTokenTypes +import com.jetbrains.python.PythonLanguage +import com.jetbrains.python.psi.* +import com.jetbrains.python.psi.impl.PyPsiUtils + +/** + * + * This data class collect information about possible DataFrame. + * @param psiName - name of PsiElement, that could be DataFrame type + * @param needValidatorCheck - boolean flag to switch on check for validating columns name + * LanguageNamesValidation.INSTANCE.forLanguage(PythonLanguage.getInstance()).isIdentifier() + * + */ +data class PandasDataFrameCandidate(val psiName: String, val needValidatorCheck: Boolean) + + +// Priority value to control order in CompletionResultSet +private const val DATAFRAME_COLUMN_PRIORITY = 100.0 + +fun getCompleteAttribute(parameters: CompletionParameters): List { + val currentElement: PsiElement? + val needValidatorCheck: Boolean + val parentLambdaExpression: PyExpression? + val candidates = mutableListOf() + val callInnerReferenceExpression = getCallInnerReferenceExpression(parameters) + + findCompleteAttribute(parameters)?.let { + currentElement = it.first + needValidatorCheck = it.second + when (currentElement) { + is PyCallExpression, is PyParenthesizedExpression -> { + return emptyList() + } + is PyExpression -> { + parentLambdaExpression = collectParentOfLambdaExpression(currentElement, callInnerReferenceExpression) + parentLambdaExpression?.text?.let { name -> candidates.add(PandasDataFrameCandidate(name, needValidatorCheck)) } + } + !is PyReferenceExpression -> { + return emptyList() + } + } + currentElement.text?.let { name -> candidates.add(PandasDataFrameCandidate(name, needValidatorCheck)) } + } + if (!candidates.isEmpty()) return candidates + return getPossibleObjectsDataFrame(parameters, callInnerReferenceExpression) + +} + +private fun getPossibleObjectsDataFrame(parameters: CompletionParameters, + callInnerReferenceExpression: PyExpression?): List { + return setOfNotNull( + callInnerReferenceExpression?.text, + getSliceSubscriptionReferenceExpression(parameters)?.text, + getAttributeReferenceExpression(parameters)?.text + ).map { PandasDataFrameCandidate(it, false) } +} + +private fun findCompleteAttribute(parameters: CompletionParameters): Pair? { + var needCheck = true + var element = parameters.position + + if (element.prevSibling?.elementType != PyTokenTypes.DOT && element.parent?.prevSibling?.elementType == PyTokenTypes.LBRACKET) { + needCheck = false + element = element.parent + } + + val exactElement = element.prevSibling?.prevSibling + when { + exactElement != null -> + return Pair(exactElement, needCheck) + + else -> return null + } +} + +private 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) + ?: PsiTreeUtil.getParentOfType(parameters.position, PySliceExpression::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 +} + +fun processDataFrameColumns(dfName: String, + columns: List, + needValidatorCheck: Boolean, + elementOnPosition: PsiElement, + project: Project, + ignoreML: Boolean): List { + val validator = LanguageNamesValidation.INSTANCE.forLanguage(PythonLanguage.getInstance()) + return columns.mapNotNull { column -> + when { + !needValidatorCheck && elementOnPosition !is PyPlainStringElement -> { + "'${StringUtil.escapeChar(column, '\'')}'" + } + !needValidatorCheck -> StringUtil.escapeChar(column, '\'') + validator.isIdentifier(column, project) -> column + else -> null + }?.let { + val lookupElement = LookupElementBuilder.create(it).withTypeText( + PyBundle.message("pandas.completion.type.text", dfName) + ) + when (ignoreML) { + true -> PrioritizedLookupElement.withPriority(lookupElement, DATAFRAME_COLUMN_PRIORITY).asMLIgnorable() + false -> PrioritizedLookupElement.withPriority(lookupElement, DATAFRAME_COLUMN_PRIORITY) + } + } + } +} + +private fun LookupElement.asMLIgnorable(): LookupElement { + return object : LookupElementDecorator(this), MLRankingIgnorable {} +}