UTBotJava/utbot-python/src/main/kotlin/org/utbot/python/code/CodeGen.kt at 066e6f924d0c861f198f28e2be5d12dc05a7ca96 · UnitTestBot/UTBotJava · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
package org.utbot.python.code

import org.utbot.framework.plugin.api.ExecutableId
import org.utbot.framework.plugin.api.UtModel
import org.utbot.framework.plugin.api.util.UtContext
import org.utbot.framework.plugin.api.util.withUtContext
import org.utbot.python.PythonMethod
import org.utbot.python.framework.api.python.PythonClassId
import org.utbot.python.framework.codegen.PythonCgLanguageAssistant
import org.utbot.python.newtyping.general.Type


object PythonCodeGenerator {
    fun generateRunFunctionCode(
        method: PythonMethod,
        methodArguments: List<UtModel>,
        directoriesForSysPath: Set<String>,
        moduleToImport: String,
        additionalModules: Set<String> = emptySet(),
        fileForOutputName: String,
        coverageDatabasePath: String,
    ): String {
        val context = UtContext(this::class.java.classLoader)
        withUtContext(context) {
            val codegen = org.utbot.python.framework.codegen.model.PythonCodeGenerator(
                PythonClassId("TopLevelFunction"),
                paramNames = emptyMap<ExecutableId, List<String>>().toMutableMap(),
                testFramework = PythonCgLanguageAssistant.getLanguageTestFrameworkManager().testFrameworks[0],
                testClassPackageName = "",
            )
            return codegen.generateFunctionCall(
                method,
                methodArguments,
                directoriesForSysPath,
                moduleToImport,
                additionalModules,
                fileForOutputName,
                coverageDatabasePath,
            )
        }
    }

    fun generateMypyCheckCode(
        method: PythonMethod,
        methodAnnotations: Map<String, Type>,
        directoriesForSysPath: Set<String>,
        moduleToImport: String,
        namesInModule: Collection<String>,
        additionalVars: String
    ): String {
        val context = UtContext(this::class.java.classLoader)
        withUtContext(context) {
            val codegen = org.utbot.python.framework.codegen.model.PythonCodeGenerator(
                PythonClassId("TopLevelFunction"),
                paramNames = emptyMap<ExecutableId, List<String>>().toMutableMap(),
                testFramework = PythonCgLanguageAssistant.getLanguageTestFrameworkManager().testFrameworks[0],
                testClassPackageName = "",
            )
            return codegen.generateMypyCheckCode(
                method,
                methodAnnotations,
                directoriesForSysPath,
                moduleToImport,
                namesInModule,
                additionalVars
            )
        }
    }
}