From 09c2fee886f1ae43571ed53b09cf780c0b03003a Mon Sep 17 00:00:00 2001 From: Phodal Huang Date: Tue, 17 Sep 2024 23:02:36 +0800 Subject: [PATCH] feat(shire-python): add basic related type support --- .../ShirePythonPsiVariableProvider.kt | 83 +++++++++++++++++++ .../shirelang/python/util/PythonPsiUtil.kt | 25 +++++- .../resources/com.phodal.shirelang.python.xml | 3 + 3 files changed, 107 insertions(+), 4 deletions(-) create mode 100644 languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/provider/ShirePythonPsiVariableProvider.kt diff --git a/languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/provider/ShirePythonPsiVariableProvider.kt b/languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/provider/ShirePythonPsiVariableProvider.kt new file mode 100644 index 00000000..89e8ac21 --- /dev/null +++ b/languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/provider/ShirePythonPsiVariableProvider.kt @@ -0,0 +1,83 @@ +package com.phodal.shirelang.python.provider + +import com.intellij.openapi.editor.Editor +import com.intellij.openapi.project.Project +import com.intellij.psi.PsiElement +import com.jetbrains.python.PythonLanguage +import com.jetbrains.python.psi.PyClass +import com.jetbrains.python.psi.PyFile +import com.jetbrains.python.psi.PyFunction +import com.phodal.shirecore.provider.variable.PsiContextVariableProvider +import com.phodal.shirecore.provider.variable.model.PsiContextVariable +import com.phodal.shirelang.python.util.PyTestUtil +import com.phodal.shirelang.python.util.PythonPsiUtil + +class ShirePythonPsiVariableProvider : PsiContextVariableProvider { + override fun resolve( + variable: PsiContextVariable, + project: Project, + editor: Editor, + psiElement: PsiElement?, + ): String { + if (psiElement?.language !is PythonLanguage) return "" + + val underTestElement = PyTestUtil.getElementForTests(project, editor) + val underTestFile = underTestElement?.containingFile as? PyFile ?: return "" + + return when (variable) { + PsiContextVariable.CURRENT_CLASS_NAME -> { + when (underTestElement) { + is PyClass -> underTestElement.name ?: "" + is PyFunction -> underTestElement.name ?: "" + else -> "" + } + } + + PsiContextVariable.CURRENT_CLASS_CODE -> { + when (underTestElement) { + is PyClass -> underTestElement.text + is PyFunction -> underTestElement.text + else -> "" + } + } + + PsiContextVariable.CURRENT_METHOD_NAME -> { + when (underTestElement) { + is PyFunction -> underTestElement.name ?: "" + else -> "" + } + } + + PsiContextVariable.CURRENT_METHOD_CODE -> { + when (underTestElement) { + is PyFunction -> underTestElement.text + else -> "" + } + } + + PsiContextVariable.RELATED_CLASSES -> { + when (underTestElement) { + is PyFunction -> { + PythonPsiUtil.findRelatedTypes(underTestElement).mapNotNull { it?.name ?: "" } + } + + else -> listOf() + }.joinToString("\n") + } + + PsiContextVariable.SIMILAR_TEST_CASE -> TODO() + PsiContextVariable.IMPORTS -> PythonPsiUtil.getImportsInFile(underTestFile) + PsiContextVariable.IS_NEED_CREATE_FILE -> TODO() + PsiContextVariable.TARGET_TEST_FILE_NAME -> { + PyTestUtil.getTestNameExample(underTestFile.virtualFile) + } + PsiContextVariable.UNDER_TEST_METHOD_CODE -> TODO() + PsiContextVariable.FRAMEWORK_CONTEXT -> TODO() + PsiContextVariable.CODE_SMELL -> TODO() + PsiContextVariable.METHOD_CALLER -> TODO() + PsiContextVariable.CALLED_METHOD -> TODO() + PsiContextVariable.SIMILAR_CODE -> TODO() + PsiContextVariable.STRUCTURE -> TODO() + } + } +} diff --git a/languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/util/PythonPsiUtil.kt b/languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/util/PythonPsiUtil.kt index 72e65091..cfa5bb0f 100644 --- a/languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/util/PythonPsiUtil.kt +++ b/languages/shire-python/src/main/kotlin/com/phodal/shirelang/python/util/PythonPsiUtil.kt @@ -1,12 +1,29 @@ package com.phodal.shirelang.python.util import com.intellij.openapi.project.Project +import com.intellij.psi.PsiFile import com.intellij.util.concurrency.annotations.RequiresReadLock import com.jetbrains.python.psi.* import com.jetbrains.python.psi.types.PyType import com.jetbrains.python.psi.types.TypeEvalContext object PythonPsiUtil { + fun getImportsInFile(file: PsiFile): String { + if (file !is PyFile) return "" + + val fromImports = file.fromImports + .map { it.text }.distinct() + .joinToString("\n") + + val imports = file.importTargets + .asSequence() + .map { it.parent.text } + .distinct() + .joinToString("\n") + + return (imports + "\n" + fromImports).trimIndent() + } + fun getClassWithoutMethod(clazz: PyClass, function: PyFunction): PyClass { val classCopy = clazz.copy() as PyClass val methods = classCopy.methods @@ -63,12 +80,12 @@ object PythonPsiUtil { } @RequiresReadLock - fun findRelatedTypes(declaration: PyFunction): List { - val context = TypeEvalContext.codeCompletion(declaration.project, declaration.containingFile) + fun findRelatedTypes(function: PyFunction): List { + val context = TypeEvalContext.codeCompletion(function.project, function.containingFile) - val resultType = declaration.getReturnStatementType(context) + val resultType = function.getReturnStatementType(context) - val parameters = declaration.parameterList.parameters + val parameters = function.parameterList.parameters val parameterTypes = parameters .filterIsInstance() diff --git a/languages/shire-python/src/main/resources/com.phodal.shirelang.python.xml b/languages/shire-python/src/main/resources/com.phodal.shirelang.python.xml index b62651ee..2b1c18d8 100644 --- a/languages/shire-python/src/main/resources/com.phodal.shirelang.python.xml +++ b/languages/shire-python/src/main/resources/com.phodal.shirelang.python.xml @@ -9,5 +9,8 @@ + +