Python plugin fixes (#1848)

* plugin fixes

* Added inner class filter for test function module

---------

Co-authored-by: Vyacheslav Tamarin <tamairn.vs@ya.ru>

Author: tochilinak

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogProcessor.kt

@@ -115,7 +115,6 @@ object PythonDialogProcessor {
                 moduleToImport,
                 UtSettings.utBotGenerationTimeoutInMillis,
                 DEFAULT_TIMEOUT_FOR_RUN_IN_MILLIS,
-                visitOnlySpecifiedSource = false,
                 cgLanguageAssistant = PythonCgLanguageAssistant,
                 pythonPath = pythonPath,
             )

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogWindow.kt

@@ -5,10 +5,8 @@ import com.intellij.openapi.ui.ComboBox
 import com.intellij.openapi.ui.DialogPanel
 import com.intellij.openapi.ui.DialogWrapper
 import com.intellij.openapi.ui.ValidationInfo
-import com.intellij.ui.ColoredListCellRenderer
 import com.intellij.ui.ContextHelpLabel
 import com.intellij.ui.JBIntSpinner
-import com.intellij.ui.SimpleTextAttributes
 import com.intellij.ui.components.Panel
 import com.intellij.ui.layout.CellBuilder
 import com.intellij.ui.layout.Row
@@ -19,7 +17,6 @@ import com.jetbrains.python.refactoring.classes.PyMemberInfoStorage
 import com.jetbrains.python.refactoring.classes.membersManager.PyMemberInfo
 import com.jetbrains.python.refactoring.classes.ui.PyMemberSelectionTable
 import org.utbot.framework.UtSettings
-import org.utbot.framework.codegen.domain.TestFramework
 import java.awt.BorderLayout
 import java.util.concurrent.TimeUnit
 import org.utbot.intellij.plugin.ui.components.TestSourceDirectoryChooser
@@ -51,8 +48,6 @@ class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.proj
     private val testFrameworks =
         ComboBox(DefaultComboBoxModel(model.cgLanguageAssistant.getLanguageTestFrameworkManager().testFrameworks.toTypedArray()))
 
-    private val visitOnlySpecifiedSource = JCheckBox("Visit only specified source")
-
     private lateinit var panel: DialogPanel
 
     init {
@@ -96,12 +91,6 @@ class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.proj
             row {
                 scrollPane(functionsTable)
             }
-            row {
-                cell {
-                    component(visitOnlySpecifiedSource)
-                    component(ContextHelpLabel.create("Find argument types only in this file."))
-                }
-            }
         }
 
         updateFunctionsTable()
@@ -140,7 +129,7 @@ class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.proj
             return globalPyFunctionsToPyMemberInfo(project, functions)
         }
         return PyMemberInfoStorage(containingClass).getClassMemberInfos(containingClass)
-            .filter { it.member is PyFunction }
+            .filter { it.member is PyFunction && fineFunction(it.member as PyFunction) }
     }
 
     private fun updateFunctionsTable() {
@@ -182,7 +171,6 @@ class PythonDialogWindow(val model: PythonTestsModel) : DialogWrapper(model.proj
         model.testFramework = testFrameworks.item
         model.timeout = TimeUnit.SECONDS.toMillis(timeoutSpinnerForTotalTimeout.number.toLong())
         model.timeoutForRun = TimeUnit.SECONDS.toMillis(timeoutSpinnerForOneRun.number.toLong())
-        model.visitOnlySpecifiedSource = visitOnlySpecifiedSource.isSelected
         model.testSourceRootPath = testSourceFolderField.text
 
         super.doOKAction()

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonLanguageAssistant.kt

@@ -59,19 +59,27 @@ object PythonLanguageAssistant : LanguageAssistant() {
             e.getData(CommonDataKeys.PSI_ELEMENT) ?: return null
         }
 
-        val containingFunction = getContainingElement<PyFunction>(element)
-        val containingClass = getContainingElement<PyClass>(element)
+        val containingClass = getContainingElement<PyClass>(element) { fineClass(it) }
+        val containingFunction: PyFunction? =
+            if (containingClass == null)
+                getContainingElement(element) { it.parent is PsiFile && fineFunction(it) }
+            else
+                getContainingElement(element) { func ->
+                    val ancestors = getAncestors(func)
+                    ancestors.dropLast(1).all { it !is PyFunction } &&
+                            ancestors.count { it is PyClass } == 1 && fineFunction(func)
+                }
 
         if (containingClass == null) {
-            val functions = file.topLevelFunctions
+            val functions = file.topLevelFunctions.filter { fineFunction(it) }
             if (functions.isEmpty())
                 return null
 
             val focusedFunction = if (functions.contains(containingFunction)) containingFunction else null
             return Targets(functions.toSet(), null, focusedFunction, file, editor)
         }
 
-        val functions = containingClass.methods
+        val functions = containingClass.methods.filter { fineFunction(it) }
         if (functions.isEmpty())
             return null
 

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonTestsModel.kt

@@ -21,7 +21,6 @@ class PythonTestsModel(
     val currentPythonModule: String,
     var timeout: Long,
     var timeoutForRun: Long,
-    var visitOnlySpecifiedSource: Boolean,
     val cgLanguageAssistant: CgLanguageAssistant,
     val pythonPath: String,
 ) : BaseTestsModel(

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/Utils.kt

@@ -5,6 +5,9 @@ import com.intellij.openapi.roots.ProjectFileIndex
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElement
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyDecorator
+import com.jetbrains.python.psi.PyFunction
 import org.utbot.python.utils.RequirementsUtils
 import kotlin.random.Random
 
@@ -19,6 +22,12 @@ inline fun <reified T : PsiElement> getContainingElement(
     return result as? T
 }
 
+fun getAncestors(element: PsiElement): List<PsiElement> =
+    if (element.parent == null)
+        listOf(element)
+    else
+        getAncestors(element.parent) + element
+
 fun getContentRoot(project: Project, file: VirtualFile): VirtualFile {
     return ProjectFileIndex.getInstance(project)
         .getContentRootForFile(file) ?: error("Source file lies outside of a module")
@@ -38,3 +47,12 @@ fun VirtualFile.isProjectSubmodule(ancestor: VirtualFile?): Boolean {
 fun checkModuleIsInstalled(pythonPath: String, moduleName: String): Boolean {
     return RequirementsUtils.requirementsAreInstalled(pythonPath, listOf(moduleName))
 }
+
+fun fineFunction(function: PyFunction): Boolean =
+    !listOf("__init__", "__new__").contains(function.name) &&
+            function.decoratorList?.decorators?.isNotEmpty() != true  // TODO: add processing of simple decorators
+            //(function.parent !is PyDecorator || (function.parent as PyDecorator).isBuiltin)
+
+fun fineClass(pyClass: PyClass): Boolean =
+    getAncestors(pyClass).dropLast(1).all { it !is PyClass && it !is PyFunction } &&
+            pyClass.methods.any { fineFunction(it) }

utbot-python/samples/easy_samples/general.py

@@ -13,6 +13,10 @@ class Dummy:
     def propagate(self):
         return [self, self]
 
+    @staticmethod
+    def abs(x):
+        return abs(x)
+
 
 def dict_f(x, a, b, c):
     y = {1: 2}

utbot-python/samples/easy_samples/sample_classes.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 
+
 class A:
     def __init__(self, val: int):
         self.description = val
@@ -19,3 +20,17 @@ class C:
 
     def inc(self):
         self.counter += 1
+
+
+class Outer:
+    class Inner:
+        a = 1
+
+        def inc(self):
+            self.a += 1
+
+    def __init__(self):
+        self.inner = Outer.Inner()
+
+    def inc1(self):
+        self.inner.inc()

utbot-python/samples/easy_samples/subtypes.py

@@ -26,12 +26,12 @@ def func_for_P(x: P) -> None:
 
 
 class R(Protocol):
-    def f(self) -> R:
+    def f(self) -> 'R':
         ...
 
 
 class RImpl:
-    def f(self) -> RImpl:
+    def f(self) -> 'RImpl':
         return self
 
 

utbot-python/src/main/kotlin/org/utbot/python/PythonEngine.kt

@@ -158,7 +158,8 @@ class PythonEngine(
         val argumentModules = argumentValues
             .flatMap { it.allContainingClassIds }
             .map { it.moduleName }
-        val localAdditionalModules = (additionalModules + argumentModules).toSet()
+            .filterNot { it.startsWith(moduleToImport) }
+        val localAdditionalModules = (additionalModules + argumentModules + moduleToImport).toSet()
 
         return PythonCodeExecutorImpl(
             methodUnderTest,

utbot-python/src/main/kotlin/org/utbot/python/newtyping/ast/test.kt

@@ -1,23 +1,18 @@
 package org.utbot.python.newtyping.ast
 
 import org.parsers.python.PythonParser
+import org.utbot.python.PythonMethodHeader
+import org.utbot.python.code.PythonCode
 
 fun main() {
     val content = """
-    def calc_to_goal_cost(trajectory, goal):
-        ""${'"'}
-            calc to goal cost with angle difference
-        ""${'"'}
-    
-        dx = goal[0] - trajectory[-1, 0]
-        dy = goal[1] - trajectory[-1, 1]
-        error_angle = math.atan2(dy, dx)
-        cost_angle = error_angle - trajectory[-1, 2]
-        cost = abs(math.atan2(math.sin(cost_angle), math.cos(cost_angle)))
-    
-        return cost
+    class A:
+        @decorator
+        def func(x):
+            return 1
     """.trimIndent()
 
     val root = PythonParser(content).Module()
+    // val y = PythonCode.findFunctionDefinition(root, PythonMethodHeader("func", "", null))
     val x = root
 }