containing classIds for PythonTree; some fixes

tochilinak · tochilinak · commit c4265979a5cf · 2022-08-11T12:03:09.000+03:00
diff --git a/utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/Api.kt b/utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/Api.kt
@@ -274,7 +274,13 @@ sealed class PythonModel(classId: PythonClassId): UtModel(classId) {
 class PythonTreeModel(
     val tree: PythonTree.PythonTreeNode,
     classId: PythonClassId,
-): PythonModel(classId)
+): PythonModel(classId) {
+    override val allContainingClassIds: Set<PythonClassId>
+        get() {
+            val children = tree.children.map { PythonTreeModel(it, it.type) }
+            return super.allContainingClassIds + children.flatMap { it.allContainingClassIds }
+        }
+}
 
 class PythonDefaultModel(
     val repr: String,
diff --git a/utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/PythonTree.kt b/utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/PythonTree.kt
@@ -3,7 +3,9 @@ package org.utbot.framework.plugin.api
 object PythonTree {
     open class PythonTreeNode(
         val type: PythonClassId,
-    )
+    ) {
+        open val children: List<PythonTreeNode> = emptyList()
+    }
 
     class PrimitiveNode(
         type: PythonClassId,
@@ -12,19 +14,31 @@ object PythonTree {
 
     class ListNode(
         val items: List<PythonTreeNode>
-    ): PythonTreeNode(PythonClassId("builtins.list"))
+    ): PythonTreeNode(PythonClassId("builtins.list")) {
+        override val children: List<PythonTreeNode>
+            get() = items
+    }
 
     class DictNode(
         val items: Map<PythonTreeNode, PythonTreeNode>
-    ): PythonTreeNode(PythonClassId("builtins.dict"))
+    ): PythonTreeNode(PythonClassId("builtins.dict")) {
+        override val children: List<PythonTreeNode>
+            get() = items.values + items.keys
+    }
 
     class SetNode(
         val items: Set<PythonTreeNode>
-    ): PythonTreeNode(PythonClassId("builtins.set"))
+    ): PythonTreeNode(PythonClassId("builtins.set")) {
+        override val children: List<PythonTreeNode>
+            get() = items.toList()
+    }
 
     class TupleNode(
         val items: List<PythonTreeNode>
-    ): PythonTreeNode(PythonClassId("builtins.tuple"))
+    ): PythonTreeNode(PythonClassId("builtins.tuple")) {
+        override val children: List<PythonTreeNode>
+            get() = items
+    }
 
     class ReduceNode(
         type: PythonClassId,
@@ -33,5 +47,8 @@ object PythonTree {
         val state: Map<String, PythonTreeNode>,
         val listitems: List<PythonTreeNode>,
         val dictitems: Map<PythonTreeNode, PythonTreeNode>,
-    ): PythonTreeNode(type)
+    ): PythonTreeNode(type)  {
+        override val children: List<PythonTreeNode>
+            get() = args + state.values + listitems + dictitems.values + dictitems.keys
+    }
 }
diff --git a/utbot-python/src/main/kotlin/org/utbot/python/PythonEngine.kt b/utbot-python/src/main/kotlin/org/utbot/python/PythonEngine.kt
@@ -55,7 +55,7 @@ class PythonEngine(
 
                 val resultAsModel = PythonTreeModel(
                     resultJSON.output,
-                    PythonClassId(resultJSON.type)
+                    resultJSON.type
                 )
                 val result = UtExecutionSuccess(resultAsModel)
 
diff --git a/utbot-python/src/main/kotlin/org/utbot/python/PythonEvaluation.kt b/utbot-python/src/main/kotlin/org/utbot/python/PythonEvaluation.kt
@@ -1,5 +1,6 @@
 package org.utbot.python
 
+import org.utbot.framework.plugin.api.PythonClassId
 import org.utbot.framework.plugin.api.PythonTree
 import org.utbot.framework.plugin.api.UtModel
 import org.utbot.python.code.KlaxonPythonTreeParser
@@ -15,7 +16,7 @@ class EvaluationSuccess(val output: OutputData, val isException: Boolean): Evalu
     operator fun component2() = isException
 }
 
-data class OutputData(val output: PythonTree.PythonTreeNode, val type: String)
+data class OutputData(val output: PythonTree.PythonTreeNode, val type: PythonClassId)
 
 object PythonEvaluation {
     fun evaluate(
diff --git a/utbot-python/src/main/resources/python_tree_serializer.py b/utbot-python/src/main/resources/python_tree_serializer.py
@@ -1,4 +1,3 @@
-import inspect
 import types
 from itertools import zip_longest
 
@@ -8,9 +7,9 @@ class _PythonTreeSerializer:
     def get_type(py_object):
         if py_object is None:
             return 'types.NoneType'
-        module = inspect.getmodule(type(py_object))
+        module = type(py_object).__module__
         return '{module}.{name}'.format(
-            module='' if module is None else module.__name__,
+            module=module,
             name=type(py_object).__name__,
         )
 

Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ class PythonEngine(`
`55`	`55`
`56`	`56`	`val resultAsModel = PythonTreeModel(`
`57`	`57`	`resultJSON.output,`
`58`		`- PythonClassId(resultJSON.type)`
	`58`	`+ resultJSON.type`
`59`	`59`	`)`
`60`	`60`	`val result = UtExecutionSuccess(resultAsModel)`
`61`	`61`