(fix class inheritance tests)

devdoomari3 · devdoomari3 · commit c63d6539415e · 2020-03-02T12:02:59.000+09:00
diff --git a/py_codegen/type_extractor/__tests__/test_classes_with_inheritance.py b/py_codegen/type_extractor/__tests__/test_classes_with_inheritance.py
@@ -1,15 +1,44 @@
 from py_codegen.test_fixtures.classes_with_inheritance import ChildClass
 
-from py_codegen.type_extractor.nodes.BaseNodeType import BaseOption
 from py_codegen.type_extractor.nodes.ClassFound import ClassFound
-from py_codegen.type_extractor.nodes.TypeOR import TypeOR
 from py_codegen.type_extractor.__tests__.utils import traverse, cleanup
 from py_codegen.type_extractor.type_extractor import TypeExtractor
 
 
 def test_classes_with_inheritance():
-    type_collector = TypeExtractor()
+    type_extractor = TypeExtractor()
 
-    type_collector.add(None)(ChildClass)
+    type_extractor.add(None)(ChildClass)
 
-    print(type_collector)
+    classes = {
+        key: traverse(value, cleanup)
+        for (key, value) in type_extractor.collected_types.items()
+        if isinstance(value, ClassFound)
+    }
+    parent_class_a = ClassFound(
+        name='ParentClassA',
+        fields={
+            'from_parent_a': int,
+        },
+    )
+    parent_class_b = ClassFound(
+        name='ParentClassB',
+        fields={
+            'from_parent_b': str,
+        },
+    )
+    assert classes == {
+        'ParentClassA': parent_class_a,
+        'ParentClassB': parent_class_b,
+        'ChildClass': ClassFound(
+            name='ChildClass',
+            fields={
+                'b': str,
+            },
+            base_classes=[
+                parent_class_a,
+                parent_class_b,
+            ],
+        )
+    }
+    print(type_extractor)
diff --git a/py_codegen/type_extractor/__tests__/utils.py b/py_codegen/type_extractor/__tests__/utils.py
@@ -22,6 +22,10 @@ def traverse(node: NodeType, func: traverse_func_type):
             key: traverse(value, func)
             for (key, value) in node.fields.items()
         }
+        class_found_node.base_classes = [
+            traverse(base_class, func)
+            for base_class in node.base_classes
+        ]
         return func(class_found_node)
     if isinstance(node, FunctionFound):
         function_found_node = copy(node)
