Changed type priorities algorithm

Author: tochilinak

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

@@ -139,6 +139,46 @@ class PythonEngine(
         }
     }
 
+    private val inf = 1000
+
+    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) }
+                }
+            }
+        }
+        return candidates.toList().sortedByDescending { it.second }.map { it.first }
+    }
+
     private fun joinAnnotations(
         argInfoCollector: ArgInfoCollector,
         methodUnderTest: PythonMethod,
@@ -149,27 +189,14 @@ class PythonEngine(
         pythonPath: String,
         fileOfMethod: String
     ): Sequence<Map<String, ClassId>> {
-        val storageMap = argInfoCollector.getPriorityStorages()
+        val storageMap = argInfoCollector.getAllStorages()
         val userAnnotations = existingAnnotations.entries.associate {
             it.key to listOf(it.value)
         }
         val annotationCombinations = storageMap.entries.associate { (name, storages) ->
-            name to storages.map { storage ->
-                when (storage) {
-                    is ArgInfoCollector.TypeStorage -> setOf(storage.name)
-                    is ArgInfoCollector.MethodStorage -> PythonTypesStorage.findTypeWithMethod(storage.name)
-                    is ArgInfoCollector.FieldStorage -> PythonTypesStorage.findTypeWithField(storage.name)
-                    is ArgInfoCollector.FunctionArgStorage -> StubFileFinder.findTypeByFunctionWithArgumentPosition(
-                        storage.name,
-                        argumentPosition = storage.index
-                    )
-                    is ArgInfoCollector.FunctionRetStorage -> StubFileFinder.findTypeByFunctionReturnValue(
-                        storage.name
-                    )
-                    else -> setOf(storage.name)
-                }
-            }.flatten().toSet().toList()
+            name to findTypeCandidates(storages)
         }
+
         return MypyAnnotations.mypyCheckAnnotations(
             methodUnderTest,
             userAnnotations + annotationCombinations,

utbot-python/src/main/kotlin/org/utbot/python/code/ArgInfoCollector.kt

@@ -46,12 +46,13 @@ class ArgInfoCollector(val method: PythonMethod, val argumentTypes: List<ClassId
         val functionRetStorages: MutableSet<FunctionRetStorage> = mutableSetOf()
     ) {
         fun toList(): List<BaseStorage> {
-            return (listOf(typeStorages) + listOf(
+            return listOf(
+                typeStorages,
                 methodStorages,
                 functionRetStorages,
                 fieldStorages,
                 functionRetStorages
-            ).sortedBy { it.size }).flatten()
+            ).flatten()
         }
     }
 
@@ -65,18 +66,9 @@ class ArgInfoCollector(val method: PythonMethod, val argumentTypes: List<ClassId
         visitor.visitFunctionDef(method.ast(), collectedValues)
     }
 
-    fun suggestBasedOnConstants(): List<Set<ClassId>> {
-        return method.arguments.mapIndexed { index, param ->
-            if (argumentTypes[index] == pythonAnyClassId)
-                collectedValues[param.name]?.typeStorages?.map { ClassId(it.name) }?.toSet() ?: emptySet()
-            else
-                setOf(argumentTypes[index])
-        }
-    }
-
     fun getConstants(): List<FuzzedConcreteValue> = visitor.constStorage.toList()
 
-    fun getPriorityStorages(): Map<String, List<BaseStorage>> {
+    fun getAllStorages(): Map<String, List<BaseStorage>> {
        return collectedValues.entries.associate { (argName, storage) ->
            argName to storage.toList()
        }
@@ -110,7 +102,6 @@ class ArgInfoCollector(val method: PythonMethod, val argumentTypes: List<ClassId
         val constStorage: MutableSet<FuzzedConcreteValue>
     ): ModifierVisitor<MutableMap<String, Storage>>() {
 
-        private val knownTypes = PythonTypesStorage.builtinTypes
         private fun <A, N> namePat(): Parser<(String) -> A, A, N> {
             val names: List<Parser<(String) -> A, A, N>> = paramNames.map { paramName ->
                 map0(refl(name(equal(paramName))), paramName)
@@ -142,14 +133,15 @@ class ArgInfoCollector(val method: PythonMethod, val argumentTypes: List<ClassId
                 "builtins.set" to refl(set(drop())),
                 "builtins.tuple" to refl(tuple(drop()))
             )
-            knownTypes.forEach { typeName ->
-                if (typeMap.containsKey(typeName))
-                    typeMap[typeName] = or(
-                        typeMap[typeName]!!,
-                        refl(functionCallWithoutPrefix(name(equal(typeName)), drop()))
+            PythonTypesStorage.builtinTypes.forEach { typeNameWithoutPrefix ->
+                val typeNameWithPrefix = "builtins.$typeNameWithoutPrefix"
+                if (typeMap.containsKey(typeNameWithPrefix))
+                    typeMap[typeNameWithPrefix] = or(
+                        typeMap[typeNameWithPrefix]!!,
+                        refl(functionCallWithoutPrefix(name(equal(typeNameWithoutPrefix)), drop()))
                     )
                 else
-                    typeMap[typeName] = refl(functionCallWithoutPrefix(name(equal(typeName)), drop()))
+                    typeMap[typeNameWithPrefix] = refl(functionCallWithoutPrefix(name(equal(typeNameWithoutPrefix)), drop()))
             }
             return typeMap.entries.fold(reject()) { acc, entry ->
                 or(acc, map0(typedExpr(entry.value), TypeStorage(entry.key)))

utbot-python/src/main/kotlin/org/utbot/python/typing/PythonTypeCollector.kt

@@ -91,7 +91,7 @@ object PythonTypesStorage {
 
     val builtinTypes: List<String>
         get() = TypesFromJSONStorage.preprocessedTypes.mapNotNull {
-            if (it.name.startsWith("builtins.")) it.name.removePrefix("builtins") else null
+            if (it.name.startsWith("builtins.")) it.name.removePrefix("builtins.") else null
         }
 
     private data class ProjectClass(

utbot-python/todo.md

@@ -13,7 +13,7 @@
 - [x] обрабатывать вызов функции вида module.func()
 - [ ] учитывать AnnAssign
 - [ ] нормально парсить комплексные числа
-- [ ] config file в mypy
+- [ ] разобраться с `nan` (`nan == nan` -> `False`)
 
 Проблемы могут быть с: `__getitem__`/`__class_getitem__`,
 `__add__/__radd__/__iadd__`