Python small fixes (#2250)

* Swap focusedFile and focusedClass

* Fix invalid string representation of float and keys in state-dict

* Fix code generation for tuple

* Fix \ in strings and coverage socket

* Fix missed module imports from stateInit

Author: tamarinvs19

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

@@ -44,7 +44,7 @@ object PythonLanguageAssistant : LanguageAssistant() {
         PythonDialogProcessor.createDialogAndGenerateTests(
             project,
             targets.pyClasses + targets.pyFunctions,
-            targets.focusedClass ?: targets.focusedFunction,
+            targets.focusedFunction ?: targets.focusedClass,
             targets.editor,
         )
     }
@@ -66,8 +66,8 @@ object PythonLanguageAssistant : LanguageAssistant() {
             val file = e.getData(CommonDataKeys.PSI_FILE) as? PyFile ?: return null
             val element = findPsiElement(file, editor) ?: return null
 
-            val allFunctions = file.topLevelFunctions.filter { fineFunction(it) }
-            val allClasses = file.topLevelClasses.filter { fineClass(it) }
+            val rootFunctions = file.topLevelFunctions.filter { fineFunction(it) }
+            val rootClasses = file.topLevelClasses.filter { fineClass(it) }
 
             val containingClass = getContainingElement<PyClass>(element) { fineClass(it) }
             val containingFunction: PyFunction? =
@@ -80,18 +80,18 @@ object PythonLanguageAssistant : LanguageAssistant() {
                                 ancestors.count { it is PyClass } == 1 && fineFunction(func)
                     }
 
-            if (allClasses.isEmpty()) {
-                return if (allFunctions.isEmpty()) {
+            if (rootClasses.isEmpty()) {
+                return if (rootFunctions.isEmpty()) {
                     null
                 } else {
-                    resultFunctions.addAll(allFunctions)
+                    resultFunctions.addAll(rootFunctions)
                     focusedFunction = containingFunction
                     Targets(resultClasses, resultFunctions, null, focusedFunction, editor)
                 }
             } else {
                 if (containingClass == null) {
-                    resultClasses.addAll(allClasses)
-                    resultFunctions.addAll(allFunctions)
+                    resultClasses.addAll(rootClasses)
+                    resultFunctions.addAll(rootFunctions)
                     focusedFunction = containingFunction
                 } else {
                     resultFunctions.addAll(containingClass.methods.filter { fineFunction(it) })

utbot-python/samples/easy_samples/general.py

@@ -206,3 +206,9 @@ def n(x, y):
 
 def list_of_list(x: List[List[InventoryItem]]):
     return x
+
+
+def make_tuple(x: list[int]):
+    if len(x) == 0:
+        return tuple()
+    return tuple(x)

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

@@ -339,7 +339,7 @@ class PythonEngine(
                     logger.info { "Cannot fuzz values for types: ${parameters.map { it.pythonTypeRepresentation() }}" }
                 }
             }
-            manager.disconnect()
+            manager.shutdown()
         }
     }
 }

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

@@ -18,6 +18,7 @@ import org.utbot.python.code.PythonCode
 import org.utbot.python.framework.api.python.PythonClassId
 import org.utbot.python.framework.api.python.PythonMethodId
 import org.utbot.python.framework.api.python.PythonModel
+import org.utbot.python.framework.api.python.PythonUtExecution
 import org.utbot.python.framework.api.python.RawPythonAnnotation
 import org.utbot.python.framework.api.python.util.pythonAnyClassId
 import org.utbot.python.framework.codegen.model.PythonCodeGenerator
@@ -153,8 +154,13 @@ object PythonTestGenerationProcessor {
 
             val importParamModules = notEmptyTests.flatMap { testSet ->
                 testSet.executions.flatMap { execution ->
-                    (execution.stateBefore.parameters + execution.stateAfter.parameters +
-                            listOf(execution.stateBefore.thisInstance, execution.stateAfter.thisInstance))
+                    val params = (execution.stateAfter.parameters + execution.stateBefore.parameters).toMutableSet()
+                    val self = mutableListOf(execution.stateBefore.thisInstance, execution.stateAfter.thisInstance)
+                    if (execution is PythonUtExecution) {
+                        params.addAll(execution.stateInit.parameters)
+                        self.add(execution.stateInit.thisInstance)
+                    }
+                    (params + self)
                         .filterNotNull()
                         .flatMap { utModel ->
                         (utModel as PythonModel).let {

utbot-python/src/main/kotlin/org/utbot/python/evaluation/PythonWorkerManager.kt

@@ -33,7 +33,7 @@ class PythonWorkerManager(
         connect()
     }
 
-    fun connect() {
+    private fun connect() {
         val processStartTime = System.currentTimeMillis()
         process = startProcess(listOf(
             pythonPath,
@@ -65,14 +65,18 @@ class PythonWorkerManager(
     fun disconnect() {
         workerSocket.close()
         process.destroy()
-        coverageReceiver.kill()
     }
 
-    fun reconnect() {
+    private fun reconnect() {
         disconnect()
         connect()
     }
 
+    fun shutdown() {
+        disconnect()
+        coverageReceiver.kill()
+    }
+
     fun runWithCoverage(
         fuzzedValues: FunctionArguments,
         additionalModulesToImport: Set<String>,

utbot-python/src/main/kotlin/org/utbot/python/evaluation/serialiation/PythonObjectParser.kt

@@ -5,7 +5,6 @@ import com.squareup.moshi.adapters.PolymorphicJsonAdapterFactory
 import com.squareup.moshi.kotlin.reflect.KotlinJsonAdapterFactory
 import org.utbot.python.framework.api.python.PythonClassId
 import org.utbot.python.framework.api.python.PythonTree
-import org.utbot.python.framework.api.python.util.pythonStrClassId
 
 object PythonObjectParser {
     private val moshi = Moshi.Builder()
@@ -149,10 +148,7 @@ fun PythonTree.PythonTreeNode.toMemoryObject(memoryDump: MemoryDump): String {
 
         is PythonTree.ReduceNode -> {
             val stateObjId = PythonTree.DictNode(this.state.entries.associate {
-                PythonTree.PrimitiveNode(
-                    pythonStrClassId,
-                    it.key
-                ) to it.value
+                PythonTree.fromString(it.key) to it.value
             }.toMutableMap())
             val argsIds = PythonTree.ListNode(this.args.withIndex().associate { it.index to it.value }.toMutableMap())
             val listItemsIds =

utbot-python/src/main/kotlin/org/utbot/python/framework/api/python/PythonTree.kt

@@ -305,9 +305,15 @@ object PythonTree {
     }
 
     fun fromFloat(value: Double): PrimitiveNode {
+        val stringValue =
+            when (value) {
+                Double.POSITIVE_INFINITY -> "float('inf')"
+                Double.NEGATIVE_INFINITY -> "-float('inf')"
+                else -> value.toString()
+            }
         return PrimitiveNode(
             pythonFloatClassId,
-            value.toString()
+            stringValue
         )
     }
 

utbot-python/src/main/kotlin/org/utbot/python/framework/api/python/util/PythonUtils.kt

@@ -18,6 +18,7 @@ fun String.toSnakeCase(): String {
 
 fun String.toPythonRepr(): String {
     val repr = this
+        .replace(Regex("((?<!\\\\)(\\\\\\\\)*)(\\\\)(?!\\\\)$"), "$1$3$3")
         .replace("\"", "\\\"")
         .replace("\\\\\"", "\\\"")
     return "\"$repr\""

utbot-python/src/main/kotlin/org/utbot/python/framework/codegen/model/constructor/visitor/CgPythonRenderer.kt

@@ -536,12 +536,16 @@ internal class CgPythonRenderer(
     }
 
     override fun visit(element: CgPythonTuple) {
-        print("(")
-        element.elements.renderSeparated()
-        if (element.elements.size == 1) {
-            print(",")
+        if (element.elements.isEmpty()) {
+            print("tuple()")
+        } else {
+            print("(")
+            element.elements.renderSeparated()
+            if (element.elements.size == 1) {
+                print(",")
+            }
+            print(")")
         }
-        print(")")
     }
 
     override fun visit(element: CgPythonSet) {