add class inheritance

devdoomari3 · devdoomari3 · commit 3460d8f1b879 · 2020-03-02T11:45:31.000+09:00
diff --git a/py_codegen/test_fixtures/class_with_methods.py b/py_codegen/test_fixtures/class_with_methods.py
diff --git a/py_codegen/test_fixtures/classes_with_inheritance.py b/py_codegen/test_fixtures/classes_with_inheritance.py
@@ -0,0 +1,10 @@
+class ParentClassA:
+    from_parent_a: int
+
+
+class ParentClassB:
+    from_parent_b: str
+
+
+class ChildClass(ParentClassA, ParentClassB):
+    b: str
diff --git a/py_codegen/type_extractor/__tests__/test_classes_with_inheritance.py b/py_codegen/type_extractor/__tests__/test_classes_with_inheritance.py
@@ -0,0 +1,15 @@
+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_collector.add(None)(ChildClass)
+
+    print(type_collector)
diff --git a/py_codegen/type_extractor/middlewares/class_found.py b/py_codegen/type_extractor/middlewares/class_found.py
@@ -30,6 +30,11 @@ def class_found_middleware(
         return duplicate
 
     _data_class = dataclass(_class)
+    base_classes = [
+        type_extractor.rawtype_to_node(base_cls)
+        for base_cls in list(_class.__bases__)
+        if base_cls is not object and base_cls is not tuple
+    ]
     argspec = inspect.getfullargspec(_data_class)
     module = inspect.getmodule(_class)
     filename = module and module.__file__
@@ -38,6 +43,7 @@ def class_found_middleware(
         name=_class.__name__,
         class_raw=_class,
         filePath=filename,
+        base_classes=base_classes,
         raw_fields=argspec.annotations,
         fields=fields,
         doc=_class.__doc__,
diff --git a/py_codegen/type_extractor/nodes/ClassFound.py b/py_codegen/type_extractor/nodes/ClassFound.py
@@ -4,7 +4,7 @@
     Dict,
     Any,
     Callable,
-    Set)
+    Set, List)
 
 from py_codegen.type_extractor.nodes.BaseNodeType import BaseNodeType, NodeType, BaseOption
 
@@ -16,6 +16,7 @@ class ClassFound(BaseNodeType):  # type: ignore
     filePath: str = field(default='')
     raw_fields: Dict[str, Any] = field(default_factory=dict)
     doc: str = field(default='')
+    base_classes: List['ClassFound'] = field(default_factory=list)
     class_raw: Optional[type] = None
     INTERNAL_fields_extra: Optional[Dict[str, Any]] = None
     options: Set[BaseOption] = field(default_factory=set)
