Support __setstate__ for custom state of python objects (#2300)

Support __setstate__ for custom state

Author: tamarinvs19

utbot-python/samples/easy_samples/setstate_test.py

@@ -0,0 +1,20 @@
+import typing
+
+
+class MyClass:
+    __slots__ = ['a', 'b']
+
+    def __init__(self, a: int, b: typing.Dict[int, str]):
+        self.a = a
+        self.b = b
+
+    def __eq__(self, other):
+        return self.a == other.a and self.b == other.b
+
+    def __setstate__(self, state):
+        self.a = state[1]['a']
+        self.b = state[1]['b']
+
+
+def make_my_class(a: int, b: typing.Dict[int, str]) -> MyClass:
+    return MyClass(a, b)

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

@@ -147,9 +147,13 @@ fun PythonTree.PythonTreeNode.toMemoryObject(memoryDump: MemoryDump): String {
         }
 
         is PythonTree.ReduceNode -> {
-            val stateObjId = PythonTree.DictNode(this.state.entries.associate {
-                PythonTree.fromString(it.key) to it.value
-            }.toMutableMap())
+            val stateObjId = if (this.customState) {
+                this.state["state"]!!
+            } else {
+                PythonTree.DictNode(this.state.entries.associate {
+                    PythonTree.fromString(it.key) to it.value
+                }.toMutableMap())
+            }
             val argsIds = PythonTree.ListNode(this.args.withIndex().associate { it.index to it.value }.toMutableMap())
             val listItemsIds =
                 PythonTree.ListNode(this.listitems.withIndex().associate { it.index to it.value }.toMutableMap())
@@ -226,7 +230,9 @@ fun MemoryObject.toPythonTree(
             }
 
             is ReduceMemoryObject -> {
-                val stateObjs = memoryDump.getById(state) as DictMemoryObject
+                val stateObjsDraft = memoryDump.getById(state)
+                val customState = stateObjsDraft !is DictMemoryObject
+
                 val arguments = memoryDump.getById(args) as ListMemoryObject
                 val listitemsObjs = memoryDump.getById(listitems) as ListMemoryObject
                 val dictitemsObjs = memoryDump.getById(dictitems) as DictMemoryObject
@@ -236,13 +242,20 @@ fun MemoryObject.toPythonTree(
                     PythonClassId(this.constructor.module, this.constructor.kind),
                     arguments.items.map { memoryDump.getById(it).toPythonTree(memoryDump, visited) },
                 )
+                prevObj.customState = customState
                 visited[this.id] = prevObj
 
-                prevObj.state = stateObjs.items.entries.associate {
-                    (memoryDump.getById(it.key).toPythonTree(memoryDump, visited) as PythonTree.PrimitiveNode)
-                        .repr.drop(1).dropLast(1) to
-                            memoryDump.getById(it.value).toPythonTree(memoryDump, visited)
-                }.toMutableMap()
+                prevObj.state = if (prevObj.customState) {
+                    val stateObjs = memoryDump.getById(state)
+                    mutableMapOf("state" to stateObjs.toPythonTree(memoryDump, visited))
+                } else {
+                    val stateObjs = memoryDump.getById(state) as DictMemoryObject
+                    stateObjs.items.entries.associate {
+                        (memoryDump.getById(it.key).toPythonTree(memoryDump, visited) as PythonTree.PrimitiveNode)
+                            .repr.drop(1).dropLast(1) to
+                                memoryDump.getById(it.value).toPythonTree(memoryDump, visited)
+                    }.toMutableMap()
+                }
                 prevObj.listitems = listitemsObjs.items.map { memoryDump.getById(it).toPythonTree(memoryDump, visited) }
                 prevObj.dictitems = dictitemsObjs.items.entries.associate {
                     memoryDump.getById(it.key).toPythonTree(memoryDump, visited) to

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

@@ -212,6 +212,7 @@ object PythonTree {
         var state: MutableMap<String, PythonTreeNode>,
         var listitems: List<PythonTreeNode>,
         var dictitems: Map<PythonTreeNode, PythonTreeNode>,
+        var customState: Boolean = false,  // if this field is true, state must have structure {'state': <PythonTreeNode>}
     ) : PythonTreeNode(id, type) {
         constructor(
             type: PythonClassId,

utbot-python/src/main/kotlin/org/utbot/python/framework/codegen/model/constructor/tree/PythonCgVariableConstructor.kt

@@ -107,8 +107,23 @@ class PythonCgVariableConstructor(cgContext: CgContext) : CgVariableConstructor(
                     keyObj to valueObj
                 }
 
-                state.forEach { (key, value) ->
-                    obj[FieldId(objectNode.type, key)] `=` value
+                if (objectNode.customState) {
+                    val setstate = state["state"]!!
+                    val methodCall = CgMethodCall(
+                        obj,
+                        PythonMethodId(
+                            obj.type as PythonClassId,
+                            "__setstate__",
+                            NormalizedPythonAnnotation(pythonNoneClassId.name),
+                            listOf(RawPythonAnnotation(setstate.type.name))
+                        ),
+                        listOf(setstate)
+                    )
+                    +methodCall
+                } else {
+                    state.forEach { (key, value) ->
+                        obj[FieldId(objectNode.type, key)] `=` value
+                    }
                 }
                 listitems.forEach {
                     val methodCall = CgMethodCall(

utbot-python/src/main/kotlin/org/utbot/python/fuzzing/provider/ReduceValueProvider.kt

@@ -59,6 +59,10 @@ object ReduceValueProvider : ValueProvider<Type, PythonFuzzedValue, PythonMethod
             .forEach {
                 // TODO: here we need to use same as .getPythonAttributeByName but without name
                 // TODO: now we do not have fields from parents
+                // val slots = type.getPythonAttributeByName(description.pythonTypeStorage, "__slots__")
+                // if (slots != null) {
+                //     TODO: here we should use only attributes from __slots__
+                //}
                 val fields = type.getPythonAttributes()
                     .filter { attr ->
                         !(attr.meta.name.startsWith("__") && attr.meta.name.endsWith("__") && attr.meta.name.length >= 4) &&

utbot-python/src/main/kotlin/org/utbot/python/utils/RequirementsUtils.kt

@@ -3,7 +3,7 @@ package org.utbot.python.utils
 object RequirementsUtils {
     val requirements: List<String> = listOf(
         "mypy==1.0.0",
-        "utbot-executor==1.4.26",
+        "utbot-executor==1.4.31",
         "utbot-mypy-runner==0.2.8",
     )