Add protocol supporting

Author: tamarinvs19

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

@@ -18,6 +18,7 @@ import org.utbot.python.framework.api.python.util.pythonAnyClassId
 import org.utbot.python.fuzzing.PythonFeedback
 import org.utbot.python.fuzzing.PythonFuzzing
 import org.utbot.python.fuzzing.PythonMethodDescription
+import org.utbot.python.newtyping.PythonTypeStorage
 import org.utbot.python.newtyping.general.Type
 import org.utbot.python.providers.PythonFuzzedMethodDescription
 import org.utbot.python.providers.defaultPythonModelProvider
@@ -39,7 +40,8 @@ class PythonEngine(
     private val fuzzedConcreteValues: List<FuzzedConcreteValue>,
     private val selectedTypeMap: Map<String, NormalizedPythonAnnotation>,
     private val timeoutForRun: Long,
-    private val initialCoveredLines: Set<Int>
+    private val initialCoveredLines: Set<Int>,
+    private val pythonTypeStorage: PythonTypeStorage? = null,
 ) {
 
     private data class JobResult(
@@ -215,7 +217,7 @@ class PythonEngine(
 
         val coveredLines = initialCoveredLines.toMutableSet()
 
-        PythonFuzzing { description, parameterValues ->
+        PythonFuzzing(pythonTypeStorage!!) { description, parameterValues ->
 
             val (thisObject, modelList) =
                 if (methodUnderTest.containingPythonClassId == null)

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

@@ -11,6 +11,7 @@ import org.utbot.python.code.ArgInfoCollector
 import org.utbot.python.framework.api.python.NormalizedPythonAnnotation
 import org.utbot.python.framework.api.python.util.pythonAnyClassId
 import org.utbot.python.newtyping.PythonTypeDescription
+import org.utbot.python.newtyping.PythonTypeStorage
 import org.utbot.python.newtyping.general.FunctionTypeCreator
 import org.utbot.python.newtyping.runmypy.RunMypy
 import org.utbot.python.typing.AnnotationFinder.findAnnotations
@@ -92,7 +93,8 @@ object PythonTestCaseGenerator {
             emptyList(),
             method.arguments.zip(args).associate { it.first.name to NormalizedPythonAnnotation((it.second.meta as PythonTypeDescription).name.toString()) },
             timeoutForRun,
-            coveredLines
+            coveredLines,
+            PythonTypeStorage.get(storage)
         )
 
         var coverageLimit = 10
@@ -192,7 +194,7 @@ object PythonTestCaseGenerator {
                     argInfoCollector.getConstants(),
                     annotations,
                     timeoutForRun,
-                    coveredLines
+                    coveredLines,
                 )
 
                 var coverageLimit = 10

utbot-python/src/main/kotlin/org/utbot/python/fuzzing/PythonApi.kt

@@ -21,6 +21,9 @@ import org.utbot.python.fuzzing.provider.StrValueProvider
 import org.utbot.python.fuzzing.provider.TupleFixSizeValueProvider
 import org.utbot.python.fuzzing.provider.TupleValueProvider
 import org.utbot.python.fuzzing.provider.UnionValueProvider
+import org.utbot.python.newtyping.PythonProtocolDescription
+import org.utbot.python.newtyping.PythonSubtypeChecker
+import org.utbot.python.newtyping.PythonTypeStorage
 import org.utbot.python.newtyping.general.Type
 
 data class PythonFuzzedConcreteValue(
@@ -64,24 +67,31 @@ fun pythonDefaultValueProviders(idGenerator: IdGenerator<Long>) = listOf(
 )
 
 class PythonFuzzing(
-    val execute: suspend (description: PythonMethodDescription, values: List<PythonTreeModel>) -> PythonFeedback
+    val pythonTypeStorage: PythonTypeStorage,
+    val execute: suspend (description: PythonMethodDescription, values: List<PythonTreeModel>) -> PythonFeedback,
 ) : Fuzzing<Type, PythonTreeModel, PythonMethodDescription, PythonFeedback> {
-    override fun generate(description: PythonMethodDescription, type: Type): Sequence<Seed<Type, PythonTreeModel>> {
+    fun generateDefault(description: PythonMethodDescription, type: Type): Sequence<Seed<Type, PythonTreeModel>> {
         val idGenerator = PythonIdGenerator()
         return pythonDefaultValueProviders(idGenerator).asSequence().flatMap { provider ->
-            try {
-                if (provider.accept(type)) {
-                    provider.generate(description, type)
-                } else {
-                    emptySequence()
-                }
-            } catch (t: Throwable) {
-                logger.error(t) { "Error occurs in value provider: $provider" }
+            if (provider.accept(type)) {
+                provider.generate(description, type)
+            } else {
                 emptySequence()
             }
         }
     }
 
+    override fun generate(description: PythonMethodDescription, type: Type): Sequence<Seed<Type, PythonTreeModel>> {
+        var providers = generateDefault(description, type)
+        if (type.meta is PythonProtocolDescription) {
+            val subtypes = pythonTypeStorage.allTypes.filter {
+                PythonSubtypeChecker.checkIfRightIsSubtypeOfLeft(type, it, pythonTypeStorage)
+            }
+            subtypes.forEach { providers += generateDefault(description, it) }
+        }
+        return providers
+    }
+
     override suspend fun handle(description: PythonMethodDescription, values: List<PythonTreeModel>): PythonFeedback {
         return execute(description, values)
     }