diff --git a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml index 5ca125f8df51..31060c06c5ec 100644 --- a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml +++ b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml @@ -486,6 +486,9 @@ + + + diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/IPythonBuiltinConstants.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/IPythonBuiltinConstants.kt new file mode 100644 index 000000000000..f14d57f5ba01 --- /dev/null +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/IPythonBuiltinConstants.kt @@ -0,0 +1,22 @@ +package com.jetbrains.python.psi.resolve + +object IPythonBuiltinConstants { + const val DISPLAY = "display" + const val GET_IPYTHON = "get_ipython" + const val IN = "In" + const val OUT = "Out" + + const val IPYTHON_PACKAGE = "IPython" + const val CORE_PACKAGE = "core" + + const val HISTORY_MANAGER = "HistoryManager" + const val OUT_HIST_DICT = "output_hist" + const val IN_HIST_DICT = "input_hist_parsed" + + const val DISPLAY_DOTTED_PATH = "IPython.core.display.display" + const val GET_IPYTHON_DOTTED_PATH = "IPython.core.getipython.get_ipython" + const val HISTORY_MANAGER_DOTTED_PATH = "IPython.core.history.HistoryManager" + + const val DOUBLE_UNDERSCORE = "__" + const val TRIPLE_UNDERSCORE = "___" +} \ No newline at end of file diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyIPythonBuiltinReferenceResolveProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyIPythonBuiltinReferenceResolveProvider.kt new file mode 100644 index 000000000000..7b39f5f97e6e --- /dev/null +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyIPythonBuiltinReferenceResolveProvider.kt @@ -0,0 +1,73 @@ +package com.jetbrains.python.psi.resolve + +import com.intellij.psi.PsiFile +import com.intellij.psi.util.QualifiedName +import com.intellij.util.PlatformUtils +import com.jetbrains.python.PythonFileType +import com.jetbrains.python.psi.PyClass +import com.jetbrains.python.psi.PyQualifiedExpression +import com.jetbrains.python.psi.impl.ResolveResultList +import com.jetbrains.python.psi.types.TypeEvalContext + +class PyIPythonBuiltinReferenceResolveProvider : PyReferenceResolveProvider { + + private val underscoreRegex = "_[0-9]*".toRegex() + + override fun resolveName(element: PyQualifiedExpression, context: TypeEvalContext): List { + if (!needToResolve(context)) return emptyList() + val name = element.referencedName ?: return emptyList() + if (name matches underscoreRegex) { + return resolveOut(element, context) + } + when (name) { + IPythonBuiltinConstants.DISPLAY -> return resolveDisplay(element) + IPythonBuiltinConstants.GET_IPYTHON -> return resolveGetIPython(element) + IPythonBuiltinConstants.IN -> return resolveIn(element, context) + IPythonBuiltinConstants.OUT -> return resolveOut(element, context) + IPythonBuiltinConstants.DOUBLE_UNDERSCORE -> return resolveOut(element, context) + IPythonBuiltinConstants.TRIPLE_UNDERSCORE -> return resolveOut(element, context) + } + return emptyList() + } + + /** + We resolve IPython built-ins in two cases: + 1) in PyCharm PRO in Jupyter files only + 2) in DataSpell in Python files and Jupyter files + */ + private fun needToResolve(context: TypeEvalContext): Boolean { + val psiFile = context.origin ?: return false + return isJupyterFile(psiFile) || (PlatformUtils.isDataSpell() && isPythonFile(psiFile)) + } + + private fun isJupyterFile(element: PsiFile): Boolean = element.virtualFile?.extension == "ipynb" + + private fun isPythonFile(element: PsiFile): Boolean = element.virtualFile?.fileType is PythonFileType + + private fun resolveDisplay(element: PyQualifiedExpression): List { + val displayPsi = resolveTopLevelMember(QualifiedName.fromDottedString(IPythonBuiltinConstants.DISPLAY_DOTTED_PATH), + fromFoothold(element)) + return ResolveResultList.to(displayPsi) + } + + private fun resolveGetIPython(element: PyQualifiedExpression): List { + val getIPythonPsi = resolveTopLevelMember(QualifiedName.fromDottedString(IPythonBuiltinConstants.GET_IPYTHON_DOTTED_PATH), + fromFoothold(element)) + return ResolveResultList.to(getIPythonPsi) + } + + private fun resolveIn(element: PyQualifiedExpression, context: TypeEvalContext): List { + val historyClass = resolveTopLevelMember(QualifiedName.fromDottedString(IPythonBuiltinConstants.HISTORY_MANAGER_DOTTED_PATH), + fromFoothold(element)) as? PyClass ?: return emptyList() + val inHistDictPsi = historyClass.findClassAttribute(IPythonBuiltinConstants.IN_HIST_DICT, false, context) + return ResolveResultList.to(inHistDictPsi) + } + + private fun resolveOut(element: PyQualifiedExpression, context: TypeEvalContext): List { + val historyClass = resolveTopLevelMember(QualifiedName.fromDottedString(IPythonBuiltinConstants.HISTORY_MANAGER_DOTTED_PATH), + fromFoothold(element)) as? PyClass ?: return emptyList() + val outHistDictPsi = historyClass.findClassAttribute(IPythonBuiltinConstants.OUT_HIST_DICT, false, context) + return ResolveResultList.to(outHistDictPsi) + } +} +