Add dataclass and pydantic support

pyproject.toml

@@ -33,7 +33,7 @@ mkdocs = "^1.1"
 mkdocstrings = "^0.10.3"
 mkdocs-material = "^4.6.3"
 mypy = "^0.770"
-pydantic = "^1.4"
+pydantic = "^1.5.1"
 pylint = { git = "https://github.com/PyCQA/pylint.git" }
 pytest = "~5.3.5"
 pytest-cov = "^2.8.1"

src/pytkdocs/loader.py

@@ -271,7 +271,7 @@ def get_module_documentation(self, node: ObjectNode, members=None) -> Module:
                 source = None
 
         root_object = Module(
-            name=name, path=path, file_path=node.file_path, docstring=inspect.getdoc(module) or "", source=source,
+            name=name, path=path, file_path=node.file_path, docstring=inspect.getdoc(module) or "", source=source
         )
 
         if members is False:
@@ -338,6 +338,22 @@ def get_class_documentation(self, node: ObjectNode, members=None) -> Class:
             elif child_node.is_property():
                 root_object.add_child(self.get_property_documentation(child_node))
 
+        # First check if this is pdyantic compataible
+        if "__fields__" in class_.__dict__:
+            root_object.properties = ["pydantic"]
+            for field_name, model_field in class_.__dict__.get("__fields__", {}).items():
+                if self.select(field_name, members):  # type: ignore
+                    child_node = ObjectNode(obj=model_field, name=field_name, parent=node)
+                    root_object.add_child(self.get_pydantic_field_documentation(child_node))
+
+        # Handle dataclasses
+        elif "__dataclass_fields__" in class_.__dict__:
+            root_object.properties = ["dataclass"]
+            for field_name, annotation in class_.__dict__.get("__annotations__", {}).items():
+                if self.select(field_name, members):  # type: ignore
+                    child_node = ObjectNode(obj=annotation, name=field_name, parent=node)
+                    root_object.add_child(self.get_annotated_dataclass_field(child_node))
+
         return root_object
 
     def get_function_documentation(self, node: ObjectNode) -> Function:
@@ -415,6 +431,49 @@ def get_property_documentation(self, node: ObjectNode) -> Attribute:
             source=source,
         )
 
+    def get_pydantic_field_documentation(self, node: ObjectNode) -> Attribute:
+        """
+        Get the documentation for a PyDantic Field
+
+        Arguments:
+            node: The node representing the Field and its parents.
+
+        Return:
+            The documented attribute object.
+        """
+        prop = node.obj
+        path = node.dotted_path
+        properties = ["field", "pydantic"]
+        if prop.required:
+            properties.append("required")
+
+        return Attribute(
+            name=node.name,
+            path=path,
+            file_path=node.file_path,
+            docstring=prop.field_info.description,
+            attr_type=prop.type_,
+            properties=properties,
+        )
+
+    def get_annotated_dataclass_field(self, node: ObjectNode) -> Attribute:
+        """
+        Get the documentation for an dataclass annotation.
+
+        Arguments:
+            node: The node representing the annotation and its parents.
+
+        Return:
+            The documented attribute object.
+        """
+        annotation: type = node.obj
+        path = node.dotted_path
+        properties = ["field"]
+
+        return Attribute(
+            name=node.name, path=path, file_path=node.file_path, attr_type=annotation, properties=properties
+        )
+
     def get_classmethod_documentation(self, node: ObjectNode) -> Method:
         """
         Get the documentation for a class-method.

tests/fixtures/dataclass.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Person:
+    """Simple dataclass for a person's information"""
+
+    name: str
+    age: int

tests/fixtures/pydantic.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel, Field
+
+
+class Person(BaseModel):
+    """Simple Pydantic Model for a person's information"""
+
+    name: str = Field("PersonA", description="The person's name")
+    age: int = Field(18, description="The person's age which must be at minimum 18")

tests/test_loader.py

@@ -94,6 +94,37 @@ def test_loading_class():
     assert obj.docstring == "The class docstring."
 
 
+def test_loading_dataclass():
+    loader = Loader()
+    obj = loader.get_object_documentation("tests.fixtures.dataclass.Person")
+    assert obj.docstring == "Simple dataclass for a person's information"
+    assert len(obj.attributes) == 2
+    name_attr = next(attr for attr in obj.attributes if attr.name == "name")
+    assert name_attr.type == str
+    age_attr = next(attr for attr in obj.attributes if attr.name == "age")
+    assert age_attr.type == int
+    assert "dataclass" in obj.properties
+
+    not_dataclass = loader.get_object_documentation("tests.fixtures.the_package.the_module.TheClass.TheNestedClass")
+    assert "dataclass" not in not_dataclass.properties
+
+
+def test_loading_pydantic_model():
+    loader = Loader()
+    obj = loader.get_object_documentation("tests.fixtures.pydantic.Person")
+    assert obj.docstring == "Simple Pydantic Model for a person's information"
+    assert "pydantic" in obj.properties
+    assert len(obj.attributes) == 2
+    name_attr = next(attr for attr in obj.attributes if attr.name == "name")
+    assert name_attr.type == str
+    assert name_attr.docstring == "The person's name"
+    assert "pydantic" in name_attr.properties
+    age_attr = next(attr for attr in obj.attributes if attr.name == "age")
+    assert age_attr.type == int
+    assert age_attr.docstring == "The person's age which must be at minimum 18"
+    assert "pydantic" in age_attr.properties
+
+
 def test_loading_nested_class():
     loader = Loader()
     obj = loader.get_object_documentation("tests.fixtures.the_package.the_module.TheClass.TheNestedClass")