feat(shire-python): add basic related type support

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()
+        }
+    }
+}

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<PyType?> {
-        val context = TypeEvalContext.codeCompletion(declaration.project, declaration.containingFile)
+    fun findRelatedTypes(function: PyFunction): List<PyType?> {
+        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<PyTypedElement>()

languages/shire-python/src/main/resources/com.phodal.shirelang.python.xml

@@ -9,5 +9,8 @@
 
         <shireAutoTesting language="Python"
                           implementationClass="com.phodal.shirelang.python.provider.ShirePythonAutoTesting"/>
+
+        <shirePsiVariableProvider language="Python"
+                                  implementationClass="com.phodal.shirelang.python.provider.ShirePythonPsiVariableProvider"/>
     </extensions>
 </idea-plugin>