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 {}
+}