Changed PythonMethod; generation of mypy check code

Author: tochilinak

utbot-python/samples/easy_samples/annotation_tests.py

@@ -8,11 +8,16 @@
 
 
 class A(Generic[XXX]):
-    self: XXX
+    self_: XXX
 
-    def f(a, b: A[int]):
+    def f(self, a, b: A[int]):
+        self.y = b
+        self_.x = b
         pass
 
+    def g(self):
+        self.x = 1
+
 
 def square(collection: Iterable[int], x: XXX):
     result = set()

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

@@ -62,7 +62,7 @@ object PythonTestCaseGenerator {
 
         // TODO: consider static and class methods
         if (method.containingPythonClassId != null) {
-            initialArgumentTypes[0] = NormalizedPythonAnnotation(method.containingPythonClassId!!.name)
+            initialArgumentTypes[0] = NormalizedPythonAnnotation(method.containingPythonClassId.name)
         }
 
         logger.debug("Collecting hints about arguments")

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

@@ -4,25 +4,30 @@ import io.github.danielnaczo.python3parser.model.stmts.compoundStmts.functionStm
 import io.github.danielnaczo.python3parser.model.mods.Module
 import org.utbot.framework.plugin.api.UtError
 import org.utbot.framework.plugin.api.UtExecution
+import org.utbot.python.code.textToModule
 import org.utbot.python.framework.api.python.PythonClassId
 import org.utbot.python.framework.api.python.util.pythonAnyClassId
+import org.utbot.python.newtyping.general.FunctionType
 import org.utbot.python.typing.MypyAnnotations
 import org.utbot.python.utils.moduleToString
 
 data class PythonArgument(val name: String, val annotation: String?)
 
-interface PythonMethod {
-    val name: String
-    val returnAnnotation: String?
-    val arguments: List<PythonArgument>
-    val moduleFilename: String
-    fun asString(): String
-    fun ast(): FunctionDef
-    val containingPythonClassId: PythonClassId?
-    fun codeLines(): List<String> = moduleToString(Module(listOf(ast().body))).split('\n')
+open class PythonMethod(
+    val name: String,
+    val returnAnnotation: String?,
+    val arguments: List<PythonArgument>,
+    val moduleFilename: String,
+    val containingPythonClassId: PythonClassId?,
+    val codeAsString: String
+) {
+    lateinit var type: FunctionType
     fun methodSignature(): String = "$name(" + arguments.joinToString(", ") {
         "${it.name}: ${it.annotation ?: pythonAnyClassId.name}"
     } + ")"
+    open val oldAst: FunctionDef by lazy {
+        textToModule(codeAsString).functionDefs.first()
+    }
 }
 
 data class PythonTestSet(

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

@@ -90,7 +90,7 @@ class ArgInfoCollector(val method: PythonMethod, private val argumentTypes: List
     private val visitor = MatchVisitor(paramNames, mutableSetOf(), GeneralStorage())
 
     init {
-        visitor.visitFunctionDef(method.ast(), collectedValues)
+        visitor.visitFunctionDef(method.oldAst, collectedValues)
     }
 
     fun getConstants(): List<FuzzedConcreteValue> = visitor.constStorage.toList()

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

@@ -25,7 +25,7 @@ class ClassInfoCollector(pyClass: PythonClass) {
             val selfName = getSelfName(method)
             if (selfName != null) {
                 val visitor = Visitor(selfName)
-                visitor.visitFunctionDef(method.ast(), storage)
+                visitor.visitFunctionDef(method.oldAst, storage)
             }
         }
         pyClass.topLevelFields.forEach { annAssign ->
@@ -36,7 +36,7 @@ class ClassInfoCollector(pyClass: PythonClass) {
     companion object {
         fun getSelfName(method: PythonMethod): String? {
             val params = method.arguments
-            if (params.isEmpty() || method.ast().decorators.any {
+            if (params.isEmpty() || method.oldAst.decorators.any {
                     listOf(
                         "staticmethod",
                         "classmethod"
@@ -46,7 +46,7 @@ class ClassInfoCollector(pyClass: PythonClass) {
         }
 
         fun isProperty(method: PythonMethod): Boolean {
-            return method.ast().decorators.any { it.name.name == "property" }
+            return method.oldAst.decorators.any { it.name.name == "property" }
         }
     }
 

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

@@ -27,6 +27,7 @@ import org.utbot.python.PythonMethod
 import org.utbot.python.framework.api.python.NormalizedPythonAnnotation
 import org.utbot.python.framework.api.python.PythonClassId
 import org.utbot.python.utils.moduleOfType
+import org.utbot.python.utils.moduleToString
 import java.util.*
 
 private val logger = KotlinLogging.logger {}
@@ -131,42 +132,37 @@ class PythonClass(
 }
 
 class PythonMethodBody(
-    private val ast: FunctionDef,
-    override val moduleFilename: String = "",
-    override val containingPythonClassId: PythonClassId? = null
-) : PythonMethod {
-    override val name: String
-        get() = ast.name.name
+    ast: FunctionDef,
+    moduleFilename: String = "",
+    containingPythonClassId: PythonClassId? = null
+) : PythonMethod(
+    name = ast.name.name,
+    returnAnnotation = annotationToString(ast.returns),
+    arguments = getArguments(ast),
+    moduleFilename,
+    containingPythonClassId,
+    codeAsString = moduleToString(Module(listOf(ast.body)))
+) {
 
-    override val returnAnnotation: String?
-        get() = annotationToString(ast.returns)
+    override val oldAst: FunctionDef = ast
 
-    // TODO: consider cases of default and keyword arguments
-    private val getParams: List<Parameter> =
-        if (ast.parameters.isPresent)
-            ast.parameters.get().params ?: emptyList()
-        else
-            emptyList()
+    companion object {
+        fun annotationToString(annotation: Optional<Expression>): String? =
+            if (annotation.isPresent) astToString(annotation.get()).trim() else null
+
+        private fun getParams(ast: FunctionDef): List<Parameter> =
+            if (ast.parameters.isPresent)
+                ast.parameters.get().params ?: emptyList()
+            else
+                emptyList()
 
-    override val arguments: List<PythonArgument>
-        get() = getParams.map { param ->
+        fun getArguments(ast: FunctionDef): List<PythonArgument> =
+            getParams(ast).map { param ->
             PythonArgument(
                 param.parameterName.name,
                 annotationToString(param.annotation)
             )
         }
-
-    override fun asString(): String {
-        return astToString(ast)
-    }
-
-    override fun ast(): FunctionDef {
-        return ast
-    }
-
-    companion object {
-        fun annotationToString(annotation: Optional<Expression>): String? =
-            if (annotation.isPresent) astToString(annotation.get()).trim() else null
     }
 }
 

utbot-python/src/main/kotlin/org/utbot/python/framework/codegen/model/PythonCodeGenerator.kt

@@ -193,13 +193,13 @@ class PythonCodeGenerator(
         }
 
         val functionPrefix = "__mypy_check"
-        val functionName = "def ${functionPrefix}_{method.name}(${parameters.joinToString(", ")}):"  // TODO: in future can be "async def"
+        val functionName = "def ${functionPrefix}_${method.name}(${parameters.joinToString(", ")}):"  // TODO: in future can be "async def"
 
         val mypyCheckCode = listOf(
             renderer.toString(),
             "",
             functionName,
-        ) + method.codeLines().map { "    $it" }
+        ) + method.codeAsString.split("\n").map { "    $it" }
         return mypyCheckCode.joinToString("\n")
     }
 }

utbot-python/src/main/kotlin/org/utbot/python/newtyping/AnnotationFromMypy.kt

@@ -305,4 +305,12 @@ class UnknownAnnotationNode: PythonAnnotationNode() {
     override fun initializeType(): Type {
         return pythonAnyType
     }
+}
+
+fun main() {
+    val json = MypyAnnotation::class.java.getResource("/annotation_sample.json")!!.readText()
+    val storage = readMypyAnnotationStorage(json)
+    val A = storage.definitions["annotation_tests"]!!["A"]!!.annotation.asUtBotType
+    val attrs = A.getPythonAttributes().map { it.name }
+    println(attrs)
 }

utbot-python/src/main/kotlin/org/utbot/python/newtyping/PythonType.kt

@@ -30,12 +30,7 @@ fun Type.isPythonObjectType(): Boolean {
 }
 
 fun Type.pythonTypeRepresentation(): String {
-    val description = pythonDescription()
-    val root = description.name.prefix.joinToString() + "." + description.name.name
-    val params = pythonAnnotationParameters()
-    if (params.isEmpty())
-        return root
-    return "$root[${params.joinToString { it.pythonTypeRepresentation() }}]"
+    return pythonDescription().getTypeRepresentation(this)
 }
 
 class PythonTypeStorage(
@@ -85,6 +80,13 @@ sealed class PythonTypeDescription(name: Name) : TypeMetaDataWithName(name) {
         getNamedMembers(type).find { it.name == name }
     open fun createTypeWithNewAnnotationParameters(like: Type, newParams: List<Type>): Type =  // overriden for Callable
         DefaultSubstitutionProvider.substituteAll(like.getOrigin(), newParams)
+    open fun getTypeRepresentation(type: Type): String {  // overriden for Callable
+        val root = name.prefix.joinToString() + "." + name.name
+        val params = getAnnotationParameters(type)
+        if (params.isEmpty())
+            return root
+        return "$root[${params.joinToString { it.pythonTypeRepresentation() }}]"
+    }
 }
 
 sealed class PythonCompositeTypeDescription(
@@ -234,6 +236,14 @@ class PythonCallableTypeDescription(
             )
         }
     }
+
+    override fun getTypeRepresentation(type: Type): String {
+        val functionType = castToCompatibleTypeApi(type)
+        val root = name.prefix.joinToString(".") + "." + name.name
+        return "$root[[${
+            functionType.arguments.joinToString(separator = ".") { it.pythonTypeRepresentation() }
+        }], ${functionType.returnValue.pythonTypeRepresentation()}]"
+    }
 }
 
 // Special Python annotations

utbot-python/src/main/kotlin/org/utbot/python/newtyping/runmypy/RunMypy.kt

@@ -1,7 +1,12 @@
 package org.utbot.python.newtyping.runmypy
 
-import org.utbot.python.newtyping.MypyAnnotationStorage
-import org.utbot.python.newtyping.readMypyAnnotationStorage
+import org.utbot.python.PythonArgument
+import org.utbot.python.PythonMethod
+import org.utbot.python.code.PythonCodeGenerator.generateMypyCheckCode
+import org.utbot.python.framework.api.python.NormalizedPythonAnnotation
+import org.utbot.python.newtyping.*
+import org.utbot.python.newtyping.general.FunctionType
+import org.utbot.python.utils.Cleaner
 import org.utbot.python.utils.TemporaryFileManager
 import org.utbot.python.utils.runCommand
 import java.io.File
@@ -71,21 +76,49 @@ private fun setConfigFile(): File {
     return file
 }
 
+fun checkSuggestedSignatureWithDMypy(
+    method: PythonMethod,
+    directoriesForSysPath: Set<String>,
+    moduleToImport: String
+): String {
+    val description = method.type.pythonDescription() as PythonCallableTypeDescription
+    val annotationMap =
+        (description.argumentNames zip method.type.arguments.map { it.pythonTypeRepresentation() }).associate {
+            Pair(it.first, NormalizedPythonAnnotation(it.second))
+        }
+    val mypyCode = generateMypyCheckCode(method, annotationMap, directoriesForSysPath, moduleToImport)
+    return mypyCode
+}
+
 fun main() {
     TemporaryFileManager.setup()
     val configFile = setConfigFile()
-    println(
-        readMypyAnnotationStorageAndInitialErrors(
-            "python3",
-            "/home/tochilinak/Documents/projects/utbot/UTBotJava/utbot-python/samples/easy_samples/annotation_tests.py",
-            configFile
-        )
-    )
+    val filePath = "/home/tochilinak/Documents/projects/utbot/UTBotJava/utbot-python/samples/easy_samples/annotation_tests.py"
+    val (storage, mypyOut) = readMypyAnnotationStorageAndInitialErrors("python3", filePath, configFile)
+    println(mypyOut)
     println(
         checkWithDMypy(
             "python3",
-            "/home/tochilinak/Documents/projects/utbot/UTBotJava/utbot-python/samples/easy_samples/annotation_tests.py",
+            filePath,
             configFile
         )
     )
+    val type = storage.definitions["annotation_tests"]!!["same_annotations"]!!.annotation.asUtBotType as FunctionType
+    val pythonMethod = PythonMethod(
+        "same_annotations",
+        type.returnValue.pythonTypeRepresentation(),
+        (type.arguments zip (type.pythonDescription() as PythonCallableTypeDescription).argumentNames).map {
+            PythonArgument(it.second, it.first.pythonTypeRepresentation())
+        },
+        filePath,
+        null,
+        "result = set()\n" +
+                    "for elem in collection:\n" +
+                    "    result.add(elem ** 2)\n" +
+                    "return result\n"
+    )
+    pythonMethod.type = type
+    println(checkSuggestedSignatureWithDMypy(pythonMethod, emptySet(), "annotation_tests"))
+
+    Cleaner.doCleaning()
 }

utbot-python/src/main/resources/requirements.txt

@@ -3,4 +3,4 @@ astor
 typeshed-client
 coverage
 utbot-executor
-utbot-mypy-runner==0.1.4
+utbot-mypy-runner==0.1.5

utbot-python/src/test/kotlin/org/utbot/python/newtyping/AnnotationFromMypyKtTest.kt

@@ -109,4 +109,11 @@ internal class AnnotationFromMypyKtTest {
             (square.pythonDescription() as PythonCallableTypeDescription).argumentNames == listOf("collection", "x")
         )
     }
+
+    @Test
+    fun testCustomClassAttributes() {
+        val A = storage.definitions["annotation_tests"]!!["A"]!!.annotation.asUtBotType
+        val attrs = A.getPythonAttributes().map { it.name }
+        assertTrue(attrs.containsAll(listOf("y", "x", "self_")))
+    }
 }