parse types from annotations

Author: tochilinak

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

@@ -22,6 +22,7 @@ import org.utbot.python.code.PythonCodeGenerator.generateTestCode
 import org.utbot.python.code.PythonCodeGenerator.saveToFile
 import org.utbot.python.PythonMethod
 import org.utbot.python.PythonTestCaseGenerator
+import org.utbot.python.code.AnnotationProcessor.getTypesFromAnnotation
 import org.utbot.python.normalizeAnnotation
 import org.utbot.python.typing.PythonTypesStorage
 
@@ -93,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,

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

@@ -249,9 +249,24 @@ fun <A, B, C, D> classMethod(
     fname: Parser<A, B, Name>,
     fattributeId: Parser<B, C, String>,
     farguments: Parser<C, D, Arguments>
-):Parser<A, D, Atom> =
+): Parser<A, D, Atom> =
     atom(refl(fname), list2(refl(attribute(fattributeId)), refl(farguments)))
 
+fun <A, B> nameWithPrefixFromAtom(
+    fname: Parser<A, B, String>
+): Parser<A, B, Atom> =
+    Parser { node, x ->
+        if (node.atomElement !is Name)
+            return@Parser Error()
+        var res = (node.atomElement as Name).id.name
+        for (elem in node.trailers) {
+            if (elem !is Attribute)
+                break
+            res += "." + elem.attr.name
+        }
+        fname.go(res, x)
+    }
+
 fun <A, B, C> functionCallWithoutPrefix(
     fname: Parser<A, B, Name>,
     farguments: Parser<B, C, Arguments>

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

@@ -4,6 +4,7 @@ import io.github.danielnaczo.python3parser.Python3Lexer
 import io.github.danielnaczo.python3parser.Python3Parser
 import io.github.danielnaczo.python3parser.model.AST
 import io.github.danielnaczo.python3parser.model.expr.Expression
+import io.github.danielnaczo.python3parser.model.expr.atoms.Atom
 import io.github.danielnaczo.python3parser.model.expr.atoms.Name
 import io.github.danielnaczo.python3parser.model.expr.atoms.Num
 import io.github.danielnaczo.python3parser.model.expr.atoms.Str
@@ -45,12 +46,7 @@ class PythonCode(private val body: Module, val filename: String? = null) {
 
     companion object {
         fun getFromString(code: String, filename: String? = null): PythonCode {
-            val lexer = Python3Lexer(fromString(code))
-            val tokens = CommonTokenStream(lexer)
-            val parser = Python3Parser(tokens)
-            val moduleVisitor = ModuleVisitor()
-            val ast = moduleVisitor.visit(parser.file_input()) as Module
-
+            val ast = textToModule(code)
             return PythonCode(ast, filename)
         }
     }
@@ -106,3 +102,36 @@ fun astToString(stmt: Statement): String {
     val modulePrettyPrintVisitor = ModulePrettyPrintVisitor()
     return modulePrettyPrintVisitor.visitModule(Module(listOf(stmt)), IndentationPrettyPrint(0))
 }
+
+fun textToModule(code: String): Module {
+    val lexer = Python3Lexer(fromString(code))
+    val tokens = CommonTokenStream(lexer)
+    val parser = Python3Parser(tokens)
+    val moduleVisitor = ModuleVisitor()
+    return moduleVisitor.visit(parser.file_input()) as Module
+}
+
+object AnnotationProcessor {
+    // get only types with modules in prefixes
+    fun getTypesFromAnnotation(annotation: String): Set<String> {
+        val annotationAST = textToModule(annotation)
+        val visitor = Visitor()
+        val result = mutableSetOf<String>()
+        visitor.visitModule(annotationAST, result)
+        return result
+    }
+
+    private class Visitor: ModifierVisitor<MutableSet<String>>() {
+        override fun visitAtom(atom: Atom, param: MutableSet<String>): AST {
+            parse(
+                nameWithPrefixFromAtom(apply()),
+                onError = null,
+                atom
+            ) { it } ?.let { typeName ->
+                param.add(typeName)
+            }
+
+            return super.visitAtom(atom, param)
+        }
+    }
+}