activated isCancelled button; restructured PythonEngine

Author: tochilinak

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

@@ -94,7 +94,7 @@ object PythonDialogProcessor {
                 val pythonPath = model.srcModule.sdk?.homePath ?: error("Couldn't find Python interpreter")
                 val testSourceRoot = model.testSourceRoot!!.path
                 val filePath = model.file.virtualFile.path
-                
+
                 // PythonCodeCollector.refreshProjectClassesList(model.project.basePath!!)
                 PythonTypesStorage.refreshProjectClassesList(
                     filePath,
@@ -114,7 +114,7 @@ object PythonDialogProcessor {
                         pythonPath,
                         model.project.basePath!!,
                         filePath
-                    )
+                    ) { indicator.isCanceled }
                 }
 
                 val tests = pythonMethods.map { method ->
@@ -125,7 +125,7 @@ object PythonDialogProcessor {
                     if (it.executions.isEmpty() && it.errors.isEmpty()) it.method.name else null
                 }
 
-                if (functionsWithoutTests.isNotEmpty()) {
+                if (functionsWithoutTests.isNotEmpty() && !indicator.isCanceled) {
                     showErrorDialogLater(
                         project,
                         message = "Cannot create tests for the following functions: " + functionsWithoutTests.joinToString { it },

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

@@ -1,6 +1,7 @@
 package org.utbot.python
 
 import org.utbot.framework.plugin.api.*
+import org.utbot.fuzzer.FuzzedConcreteValue
 import org.utbot.fuzzer.FuzzedMethodDescription
 import org.utbot.fuzzer.fuzz
 import org.utbot.fuzzer.names.MethodBasedNameSuggester
@@ -20,190 +21,79 @@ class PythonEngine(
     private val directoriesForSysPath: List<String>,
     private val moduleToImport: String,
     private val pythonPath: String,
-    private val projectRoot: String,
-    private val fileOfMethod: String
+    private val fuzzedConcreteValues: List<FuzzedConcreteValue>,
+    private val selectedTypeMap: Map<String, ClassId>
 ) {
     fun fuzzing(): Sequence<PythonResult> = sequence {
-        val argumentTypes = methodUnderTest.arguments.map {
-            annotationToClassId(
-                it.annotation,
-                pythonPath,
-                projectRoot,
-                fileOfMethod,
-                directoriesForSysPath,
-                testSourceRoot
-            )
-        }
-        val existingAnnotations = mutableMapOf<String, String>()
-        argumentTypes.forEachIndexed { index, classId ->
-            if (classId != pythonAnyClassId)
-                existingAnnotations[methodUnderTest.arguments[index].name] = classId.name
+        val types = methodUnderTest.arguments.map {
+            selectedTypeMap[it.name] ?: pythonAnyClassId
         }
 
-        val argInfoCollector = ArgInfoCollector(methodUnderTest, argumentTypes)
         val methodUnderTestDescription = FuzzedMethodDescription(
             methodUnderTest.name,
             pythonAnyClassId,
-            argumentTypes,
-            argInfoCollector.getConstants()
+            types,
+            fuzzedConcreteValues
         ).apply {
             compilableName = methodUnderTest.name // what's the difference with ordinary name?
             parameterNameMap = { index -> methodUnderTest.arguments.getOrNull(index)?.name }
         }
 
-        val annotations: Sequence<Map<String, ClassId>> =
-            if (existingAnnotations.size == methodUnderTest.arguments.size)
-                sequenceOf(
-                    existingAnnotations.mapValues { entry -> ClassId(entry.value) }
-                )
-            else
-                joinAnnotations(
-                    argInfoCollector,
-                    methodUnderTest,
-                    existingAnnotations,
-                    testSourceRoot,
-                    moduleToImport,
-                    directoriesForSysPath,
-                    pythonPath,
-                    fileOfMethod
-                )
-
-        // model provider argwith fallback?
-        // attempts?
-
         var testsGenerated = 0
+        fuzz(methodUnderTestDescription, concreteTypesModelProvider).forEach { values ->
+            val modelList = values.map { it.model }
+            val evalResult = PythonEvaluation.evaluate(
+                methodUnderTest,
+                modelList,
+                testSourceRoot,
+                directoriesForSysPath,
+                moduleToImport,
+                pythonPath
+            )
+            if (evalResult is EvaluationError)
+                return@sequence
 
-        annotations.forEach { types ->
-            val classIds = methodUnderTest.arguments.map {
-                types[it.name] ?: pythonAnyClassId
-            }
-            val substitutedDescription = substituteTypesByIndex(methodUnderTestDescription, classIds)
-            fuzz(substitutedDescription, concreteTypesModelProvider).forEach { values ->
-                val modelList = values.map { it.model }
-
-                // execute method to get function return
-                // what if exception happens?
-                val evalResult = PythonEvaluation.evaluate(
-                    methodUnderTest,
-                    modelList,
-                    testSourceRoot,
-                    directoriesForSysPath,
-                    moduleToImport,
-                    pythonPath
-                )
-                if (evalResult is EvaluationError)
-                    return@sequence
-
-                val (resultJSON, isException) = evalResult as EvaluationSuccess
-
-                if (isException) {
-                    yield(PythonError(UtError(resultJSON.output, Throwable()), modelList))
-                } else {
-
-                    // some types cannot be used as return types in tests (like socket or memoryview)
-                    val outputType = ClassId(resultJSON.type)
-                    if (PythonTypesStorage.getTypeByName(outputType)?.returnRenderType == ReturnRenderType.NONE)
-                        return@sequence
-
-                    val resultAsModel = PythonDefaultModel(resultJSON.output, "")
-                    val result = UtExecutionSuccess(resultAsModel)
+            val (resultJSON, isException) = evalResult as EvaluationSuccess
 
-                    val nameSuggester = sequenceOf(ModelBasedNameSuggester(), MethodBasedNameSuggester())
-                    val testMethodName = try {
-                        nameSuggester.flatMap { it.suggest(methodUnderTestDescription, values, result) }.firstOrNull()
-                    } catch (t: Throwable) {
-                        null
-                    }
+            if (isException) {
+                yield(PythonError(UtError(resultJSON.output, Throwable()), modelList))
+            } else {
 
-                    yield(
-                        PythonExecution(
-                            UtExecution(
-                                stateBefore = EnvironmentModels(null, modelList, emptyMap()),
-                                stateAfter = EnvironmentModels(null, modelList, emptyMap()),
-                                result = result,
-                                instrumentation = emptyList(),
-                                path = mutableListOf(), // ??
-                                fullPath = emptyList(), // ??
-                                testMethodName = testMethodName?.testName,
-                                displayName = testMethodName?.displayName
-                            ),
-                            modelList
-                        )
-                    )
-                }
-
-                testsGenerated += 1
-                if (testsGenerated == 100)
+                // some types cannot be used as return types in tests (like socket or memoryview)
+                val outputType = ClassId(resultJSON.type)
+                if (PythonTypesStorage.getTypeByName(outputType)?.returnRenderType == ReturnRenderType.NONE)
                     return@sequence
-            }
-        }
-    }
 
-    private val inf = 1000
+                val resultAsModel = PythonDefaultModel(resultJSON.output, "")
+                val result = UtExecutionSuccess(resultAsModel)
 
-    private fun increaseValue(map: MutableMap<String, Int>, key: String) {
-        if (map[key] == inf)
-            return
-        map[key] = (map[key] ?: 0) + 1
-    }
-
-    private fun findTypeCandidates(
-        storages: List<ArgInfoCollector.BaseStorage>?
-    ): List<String> {
-        val candidates = mutableMapOf<String, Int>() // key: type, value: priority
-        PythonTypesStorage.builtinTypes.associateByTo(destination = candidates, { it }, { 0 })
-        storages?.forEach { argInfoStorage ->
-            when (argInfoStorage) {
-                is ArgInfoCollector.TypeStorage -> candidates[argInfoStorage.name] = inf
-                is ArgInfoCollector.MethodStorage -> {
-                    val typesWithMethod = PythonTypesStorage.findTypeWithMethod(argInfoStorage.name)
-                    typesWithMethod.forEach { increaseValue(candidates, it) }
-                }
-                is ArgInfoCollector.FieldStorage -> {
-                    val typesWithField = PythonTypesStorage.findTypeWithField(argInfoStorage.name)
-                    typesWithField.forEach { increaseValue(candidates, it) }
-                }
-                is ArgInfoCollector.FunctionArgStorage -> {
-                    StubFileFinder.findTypeByFunctionWithArgumentPosition(
-                        argInfoStorage.name,
-                        argumentPosition = argInfoStorage.index
-                    ).forEach { increaseValue(candidates, it) }
-                }
-                is ArgInfoCollector.FunctionRetStorage -> {
-                    StubFileFinder.findTypeByFunctionReturnValue(
-                        argInfoStorage.name
-                    ).forEach { increaseValue(candidates, it) }
+                val nameSuggester = sequenceOf(ModelBasedNameSuggester(), MethodBasedNameSuggester())
+                val testMethodName = try {
+                    nameSuggester.flatMap { it.suggest(methodUnderTestDescription, values, result) }.firstOrNull()
+                } catch (t: Throwable) {
+                    null
                 }
+
+                yield(
+                    PythonExecution(
+                        UtExecution(
+                            stateBefore = EnvironmentModels(null, modelList, emptyMap()),
+                            stateAfter = EnvironmentModels(null, modelList, emptyMap()),
+                            result = result,
+                            instrumentation = emptyList(),
+                            path = mutableListOf(), // ??
+                            fullPath = emptyList(), // ??
+                            testMethodName = testMethodName?.testName,
+                            displayName = testMethodName?.displayName
+                        ),
+                        modelList
+                    )
+                )
             }
-        }
-        return candidates.toList().sortedByDescending { it.second }.map { it.first }
-    }
 
-    private fun joinAnnotations(
-        argInfoCollector: ArgInfoCollector,
-        methodUnderTest: PythonMethod,
-        existingAnnotations: Map<String, String>,
-        testSourceRoot: String,
-        moduleToImport: String,
-        directoriesForSysPath: List<String>,
-        pythonPath: String,
-        fileOfMethod: String
-    ): Sequence<Map<String, ClassId>> {
-        val storageMap = argInfoCollector.getAllStorages()
-        val userAnnotations = existingAnnotations.entries.associate {
-            it.key to listOf(it.value)
+            testsGenerated += 1
+            if (testsGenerated == 100)
+                return@sequence
         }
-        val annotationCombinations = storageMap.entries.associate { (name, storages) ->
-            name to findTypeCandidates(storages)
-        }
-
-        return MypyAnnotations.mypyCheckAnnotations(
-            methodUnderTest,
-            userAnnotations + annotationCombinations,
-            testSourceRoot,
-            moduleToImport,
-            directoriesForSysPath + listOf(File(fileOfMethod).parentFile.path),
-            pythonPath
-        )
     }
 }