Fixed Protocol serialization

tochilinak · tochilinak · commit 2eb02bdb3565 · 2022-09-23T12:29:25.000+03:00
diff --git a/utbot-python/src/main/kotlin/org/utbot/python/typing/PythonAnnotations.kt b/utbot-python/src/main/kotlin/org/utbot/python/typing/PythonAnnotations.kt
@@ -80,7 +80,7 @@ open class ConcreteAnnotation(
 }
 
 open class Protocol(
-    val methods: List<String>,
+    val members: List<String>,
     module: String,
     simpleName: String,
     type: AnnotationType = AnnotationType.Protocol,
@@ -177,7 +177,7 @@ fun main() {
     println(obj3.annotations[0].normalizedRepr)
 
     val test4 = """
-        {"nodeStorage": {"140501235228384": {"type": "Concrete", "module": "builtins", "simpleName": "int"}, "140501242115552": {"type": "NoneType"}, "140501250114464": {"type": "Union"}}, "annotations": [{"nodeId": "140501250114464", "args": [{"nodeId": "140501235228384"}, {"nodeId": "140501242115552"}]}]}
+        {"nodeStorage": {"139974471573216": {"type": "Concrete", "module": "builtins", "simpleName": "int"}, "139974471470816": {"type": "Protocol", "module": "typing", "simpleName": "Iterable", "members": ["__iter__"]}, "139974486047872": {"type": "Any"}}, "annotations": [{"nodeId": "139974471470816", "args": [{"nodeId": "139974471573216"}]}, {"nodeId": "139974486047872"}]}
     """.trimIndent()
     val obj4 = jsonAdapter.fromJson(test4)!!
     println(obj4)
diff --git a/utbot-python/src/main/resources/new_normalize_annotation_from_project.py b/utbot-python/src/main/resources/new_normalize_annotation_from_project.py
@@ -26,21 +26,37 @@ def __hash__(self):
         return hash(self.id_)
 
 
-class CustomAnnotationNode(AnnotationNode):
+class ConcreteAnnotationNode(AnnotationNode):
     def __init__(self, symbol_node):
         if isinstance(symbol_node, mypy.nodes.TypeInfo):
-            if symbol_node.is_protocol:
-                annotation_type = "Protocol"
-            else:
-                annotation_type = "Concrete"
-            super().__init__(annotation_type, id(symbol_node))
+            assert not symbol_node.is_protocol
+            super().__init__("Concrete", id(symbol_node))
             self.module = symbol_node.module_name
             self.simple_name = symbol_node._fullname.split('.')[-1]
         else:
             assert False, "Some SymbolNode wasn't considered"
 
     def encode(self):
-        return dict(super().encode(), **{"module": self.module, "simpleName": self.simple_name})
+        superclass_dict = super().encode()
+        subclass_dict = {"module": self.module, "simpleName": self.simple_name}
+        return dict(superclass_dict, **subclass_dict)
+
+
+class ProtocolAnnotationNode(AnnotationNode):
+    def __init__(self, symbol_node):
+        if isinstance(symbol_node, mypy.nodes.TypeInfo):
+            assert symbol_node.is_protocol
+            super().__init__("Protocol", id(symbol_node))
+            self.module = symbol_node.module_name
+            self.simple_name = symbol_node._fullname.split('.')[-1]
+            self.members = symbol_node.protocol_members
+        else:
+            assert False, "Some SymbolNode wasn't considered"
+
+    def encode(self):
+        superclass_dict = super().encode()
+        subclass_dict = {"module": self.module, "simpleName": self.simple_name, "members": self.members}
+        return dict(superclass_dict, **subclass_dict)
 
 
 class Annotation:
@@ -68,8 +84,10 @@ def get_annotation(mypy_type) -> Annotation:
             cur_node = AnnotationNode("List", id(mypy_type.type))
         elif mypy_type.type._fullname == "builtins.dict":
             cur_node = AnnotationNode("Dict", id(mypy_type.type))
+        elif mypy_type.type.is_protocol:
+            cur_node = ProtocolAnnotationNode(mypy_type.type)  # mypy_type.type: mypy.nodes.TypeInfo
         else:
-            cur_node = CustomAnnotationNode(mypy_type.type)  # mypy_type.type: mypy.nodes.TypeInfo
+            cur_node = ConcreteAnnotationNode(mypy_type.type)  # mypy_type.type: mypy.nodes.TypeInfo
 
         annotation_node_set.add(cur_node)
         return Annotation(cur_node.id_, children)
