Repository: Fatal1ty/mashumaro

Issue:
Take description from docstring
**Is your feature request related to a problem? Please describe.**
I would like my dataclasses to be as concise and readable as possible. This makes them easier to maintain, especially for new/inexperienced developers.

**Describe the solution you'd like**
When you add a docstring to a class attribute, use that in the description field of the generated JSON schema. For example:

```py
@dataclass
class SimpleRadiationConfig:
    Q0: float = 100
    """Fixed net radiation."""

print(build_json_schema(SimpleRadiationConfig).to_json())
{
    "type": "object",
    "title": "SimpleRadiationConfig",
    "properties": {
        "Q0": {
            "type": "number",
            "description": "Fixed net radiation.",
            "default": 100
        }
    },
    "additionalProperties": false
}
```

**Describe alternatives you've considered**
I saw #125, which achieves the same thing, but it requires the annotating all attributes as "fields".
I believe pydantic also supports this, but it requires marking all classes as pydantic.BaseModel, which feels more invasive.

**Additional context**
-

Hints from issue comments:
@Fatal1ty: hi @Peter9192

I was thinking about it from the start but the inability to distinguish an automatically added _technical_ docstring from an explicit one stopped me. Now I think that we can match the automatically generated docstring by its pattern. However, it’s worth adding a new builder parameter with the following possible values:
```python
class DocStringDocumentation(StrEnum):
    FULL = "full"  # all docstrings will be used
    EXPLICIT_ONLY = "explicit_only"  # only explicitly added, will be by default
    NONE = "none"  # none of them
```

What do you think? You can help with the naming to speed up the work.

Edit: All this applies to the dataclass documentation but not to a certain field. If you know if pydantic adds a field documentation based on the docstring, please give me more info. All I know is there is no way to document a certain field with docstring in Python.@Peter9192: Hi @Fatal1ty, thanks for the quick response! I wasn't aware of auto-generated docstrings. Can you clarify what you mean? I did find some docstring generators, but I believe you're referring to something else. Perhaps the options could be `parse_docstrings` with options `all`, `none`, and `explicit_only`. I think "explicit_only" is quite clear, can't think of better alternatives.

> All this applies to the dataclass documentation but not to a certain field.

Just to be sure: by "dataclass documentation", do you mean only the top level docstring on the dataclass? I was hoping this would be possible also for fields, i.e. class variables with a type annotation (not those explicitly defined with the [field](https://docs.python.org/3/library/dataclasses.html#dataclasses.field) function.

I may have been a bit too quick to conclude that pydantic supports this. However, I did find [a recent PR](https://github.com/pydantic/pydantic/pull/6563) that seems to add this functionality.

In the past I generated automatic API docs with [autodoc](https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#directive-option-automodule-members), which led me to believe it should be possible to extract this info quite easily. However, it seems this doesn't discriminate fields from other class members, which was okay for my use case but may be too limiting for a generic implementation.

I believe mkdocstrings also parses field docstrings, see https://github.com/mkdocstrings/python/issues/58@Fatal1ty: > I wasn't aware of auto-generated docstrings. Can you clarify what you mean?

Sure, here it is:

```python
from dataclasses import dataclass


@dataclass
class SimpleRadiationConfig:
    Q0: float = 100


print(SimpleRadiationConfig.__doc__)  # SimpleRadiationConfig(Q0: float = 100)
```

> Just to be sure: by "dataclass documentation", do you mean only the top level docstring on the dataclass?

Yes, I mean the top level docstring because it's easy to get it from `__doc__` attribute.

> I may have been a bit too quick to conclude that pydantic supports this. However, I did find https://github.com/pydantic/pydantic/pull/6563 that seems to add this functionality.

I see. They use `ast` module to parse the dataclass code. I'm not sure it's a good idea to invent a way to set field documentation in such non-standard ways that require parsing the code. I'm more inclined to use `typing.Doc` from [PEP 727](https://peps.python.org/pep-0727/). It's not accepted so far but it's already included in typing-extensions.

```python
from typing import Annotated, Doc

class User:
    name: Annotated[str, Doc("The user's name")]
    age: Annotated[int, Doc("The user's age")]
```

On the other hand, as you well noted, other tools use the comment after the field as documentation for it. It might make sense to come up with a way to connect plugins to JSON Schema generation, one of which would be to add documentation to comment-based fields.@mishamsk: just to chime in here - we use data classes for app configuration. Fields have docstrings that are then used to generate documentation and JSON schema. The latter drives a configuration UI (we are working on generating the UI automatically from the schema).

The usage of docstrings is convenient, because the text is shown in the IDE (developer-friendly), so there is no need to separately copy-paste the same or similar documentation string as the field metadata or annotation.

Yes, this means parsing the AST. In our case it is a build-time thing, so I am not at all concerned. If you'd add this to the JSON schema generator, we'll be able to use the built-in code instead of maintaining our own.

However, I can see a potential trap of mashumaro becoming bloated as undoubtedly more and more configuration will be needed to accommodate downstream usage.@Fatal1ty: Good news, everyone! I'm going to add a plugin system to the JSON schema building process. Perhaps one of the plugins included in mashumaro will be a plugin for adding descriptions from the docstring.@Fatal1ty: POC for the dataclass attribute docstring descriptions:
```python
import ast
import inspect
from dataclasses import dataclass, is_dataclass
from inspect import cleandoc
from typing import Optional

from mashumaro.jsonschema import build_json_schema
from mashumaro.jsonschema.models import Context, JSONSchema
from mashumaro.jsonschema.plugins import BasePlugin
from mashumaro.jsonschema.schema import Instance


@dataclass
class MyClass:
    x: int
    "this is x"


class ClassAttributeDocstringDescriptionPlugin(BasePlugin):
    def get_schema(
        self,
        instance: Instance,
        ctx: Context,
        schema: Optional[JSONSchema] = None,
    ) -> Optional[JSONSchema]:
        if not is_dataclass(instance.type):
            return
        try:
            source = inspect.getsource(instance.type)
        except TypeError:
            return
        tree = ast.parse(source)
        cls_def = None
        attr_name = None
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                cls_def = node
                break
        if not cls_def:
            return
        for node in cls_def.body:
            if isinstance(node, ast.AnnAssign):
                if isinstance(node.target, ast.Name):
                    attr_name = node.target.id
                next_node = cls_def.body[cls_def.body.index(node) + 1]
                if isinstance(next_node, ast.Expr) and isinstance(
                    next_node.value, ast.Constant
                ):
                    attr_doc = cleandoc(next_node.value.value)
                    if attr_doc:
                        schema.properties[attr_name].description = attr_doc
        return schema


schema = build_json_schema(
    MyClass, plugins=[ClassAttributeDocstringDescriptionPlugin()]
)
assert schema.properties["x"].description == "this is x"
```

Solution patch:
diff --git a/README.md b/README.md
index b8ff8b5..73fe03f 100644
--- a/README.md
+++ b/README.md
@@ -106,6 +106,7 @@ Table of contents
 * [JSON Schema](#json-schema)
     * [Building JSON Schema](#building-json-schema)
     * [JSON Schema constraints](#json-schema-constraints)
+    * [JSON Schema plugins](#json-schema-plugins)
     * [Extending JSON Schema](#extending-json-schema)
     * [JSON Schema and custom serialization methods](#json-schema-and-custom-serialization-methods)
 
@@ -3283,6 +3284,129 @@ Object constraints:
 * [`MinProperties`](https://json-schema.org/draft/2020-12/json-schema-validation.html#name-minproperties)
 * [`DependentRequired`](https://json-schema.org/draft/2020-12/json-schema-validation.html#name-dependentrequired)
 
+### JSON Schema plugins
+
+If the built-in functionality doesn't meet your needs, you can customize the JSON Schema generation or add support for additional types using plugins. The `mashumaro.jsonschema.plugins.BasePlugin` class provides a `get_schema` method that you can override to implement custom behavior.
+
+The plugin system works by iterating through all registered plugins and calling their `get_schema` methods. If a plugin's `get_schema` method raises a `NotImplementedError` or returns `None`, it indicates that the plugin doesn't provide the required functionality for that particular case.
+
+You can apply multiple plugins sequentially, allowing each to modify the schema in turn. This approach enables a step-by-step transformation of the schema, with each plugin contributing its specific modifications.
+
+Plugins can be registered using the `plugins` argument in either the `build_json_schema` function or the `JSONSchemaBuilder` class.
+
+The `mashumaro.jsonschema.plugins` module contains several built-in plugins. Currently, one of these plugins adds descriptions to JSON schemas using docstrings from dataclasses:
+
+```python
+from dataclasses import dataclass
+
+from mashumaro.jsonschema import build_json_schema
+from mashumaro.jsonschema.plugins import DocstringDescriptionPlugin
+
+
+@dataclass
+class MyClass:
+    """My class"""
+
+    x: int
+
+
+schema = build_json_schema(MyClass, plugins=[DocstringDescriptionPlugin()])
+print(schema.to_json())
+```
+
+<details>
+<summary>Click to show the result</summary>
+
+```json
+{
+    "type": "object",
+    "title": "MyClass",
+    "description": "My class",
+    "properties": {
+        "x": {
+            "type": "integer"
+        }
+    },
+    "additionalProperties": false,
+    "required": [
+        "x"
+    ]
+}
+```
+</details>
+
+Creating your own custom plugin is straightforward. For instance, if you want to add support for Pydantic models, you could write a plugin similar to the following:
+
+```python
+from dataclasses import dataclass
+
+from pydantic import BaseModel
+
+from mashumaro.jsonschema import build_json_schema
+from mashumaro.jsonschema.models import Context, JSONSchema
+from mashumaro.jsonschema.plugins import BasePlugin
+from mashumaro.jsonschema.schema import Instance
+
+
+class PydanticSchemaPlugin(BasePlugin):
+    def get_schema(
+        self,
+        instance: Instance,
+        ctx: Context,
+        schema: JSONSchema | None = None,
+    ) -> JSONSchema | None:
+        try:
+            if issubclass(instance.type, BaseModel):
+                pydantic_schema = instance.type.model_json_schema()
+                return JSONSchema.from_dict(pydantic_schema)
+        except TypeError:
+            return None
+
+
+class MyPydanticClass(BaseModel):
+    x: int
+
+
+@dataclass
+class MyDataClass:
+    y: MyPydanticClass
+
+
+schema = build_json_schema(MyDataClass, plugins=[PydanticSchemaPlugin()])
+print(schema.to_json())
+```
+
+<details>
+<summary>Click to show the result</summary>
+
+```json
+{
+    "type": "object",
+    "title": "MyDataClass",
+    "properties": {
+        "y": {
+            "type": "object",
+            "title": "MyPydanticClass",
+            "properties": {
+                "x": {
+                    "type": "integer",
+                    "title": "X"
+                }
+            },
+            "required": [
+                "x"
+            ]
+        }
+    },
+    "additionalProperties": false,
+    "required": [
+        "y"
+    ]
+}
+```
+</details>
+
+
 ### Extending JSON Schema
 
 Using a `Config` class it is possible to override some parts of the schema.
diff --git a/mashumaro/jsonschema/builder.py b/mashumaro/jsonschema/builder.py
index 23a31cb..28a23ab 100644
--- a/mashumaro/jsonschema/builder.py
+++ b/mashumaro/jsonschema/builder.py
@@ -1,8 +1,10 @@
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any, Optional, Type
 
 from mashumaro.jsonschema.dialects import DRAFT_2020_12, JSONSchemaDialect
 from mashumaro.jsonschema.models import Context, JSONSchema
+from mashumaro.jsonschema.plugins import BasePlugin
 from mashumaro.jsonschema.schema import Instance, get_schema
 
 try:
@@ -21,6 +23,7 @@ def build_json_schema(
     with_dialect_uri: bool = False,
     dialect: Optional[JSONSchemaDialect] = None,
     ref_prefix: Optional[str] = None,
+    plugins: Sequence[BasePlugin] = (),
 ) -> JSONSchema:
     if context is None:
         context = Context()
@@ -30,6 +33,7 @@ def build_json_schema(
             definitions=context.definitions,
             all_refs=context.all_refs,
             ref_prefix=context.ref_prefix,
+            plugins=context.plugins,
         )
     if dialect is not None:
         context.dialect = dialect
@@ -41,6 +45,8 @@ def build_json_schema(
         context.ref_prefix = ref_prefix.rstrip("/")
     elif context.ref_prefix is None:
         context.ref_prefix = context.dialect.definitions_root_pointer
+    if plugins:
+        context.plugins = plugins
     instance = Instance(instance_type)
     schema = get_schema(instance, context, with_dialect_uri=with_dialect_uri)
     if with_definitions and context.definitions:
@@ -64,6 +70,7 @@ def __init__(
         dialect: JSONSchemaDialect = DRAFT_2020_12,
         all_refs: Optional[bool] = None,
         ref_prefix: Optional[str] = None,
+        plugins: Sequence[BasePlugin] = (),
     ):
         if all_refs is None:
             all_refs = dialect.all_refs
@@ -73,6 +80,7 @@ def __init__(
             dialect=dialect,
             all_refs=all_refs,
             ref_prefix=ref_prefix.rstrip("/"),
+            plugins=plugins,
         )
 
     def build(self, instance_type: Type) -> JSONSchema:
diff --git a/mashumaro/jsonschema/models.py b/mashumaro/jsonschema/models.py
index 3d0f785..b194ccc 100644
--- a/mashumaro/jsonschema/models.py
+++ b/mashumaro/jsonschema/models.py
@@ -1,15 +1,21 @@
 import datetime
 import ipaddress
+from collections.abc import Sequence
 from dataclasses import MISSING, dataclass, field
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
-from typing_extensions import Self, TypeAlias
+from typing_extensions import TYPE_CHECKING, Self, TypeAlias
 
 from mashumaro.config import BaseConfig
 from mashumaro.helper import pass_through
 from mashumaro.jsonschema.dialects import DRAFT_2020_12, JSONSchemaDialect
 
+if TYPE_CHECKING:  # pragma: no cover
+    from mashumaro.jsonschema.plugins import BasePlugin
+else:
+    BasePlugin = Any
+
 try:
     from mashumaro.mixins.orjson import (
         DataClassORJSONMixin as DataClassJSONMixin,
@@ -190,3 +196,4 @@ class Context:
     definitions: dict[str, JSONSchema] = field(default_factory=dict)
     all_refs: Optional[bool] = None
     ref_prefix: Optional[str] = None
+    plugins: Sequence[BasePlugin] = ()
diff --git a/mashumaro/jsonschema/plugins.py b/mashumaro/jsonschema/plugins.py
new file mode 100644
index 0000000..3396c9c
--- /dev/null
+++ b/mashumaro/jsonschema/plugins.py
@@ -0,0 +1,28 @@
+from dataclasses import is_dataclass
+from inspect import cleandoc
+from typing import Optional
+
+from mashumaro.jsonschema.models import Context, JSONSchema
+from mashumaro.jsonschema.schema import Instance
+
+
+class BasePlugin:
+    def get_schema(
+        self,
+        instance: Instance,
+        ctx: Context,
+        schema: Optional[JSONSchema] = None,
+    ) -> Optional[JSONSchema]:
+        pass
+
+
+class DocstringDescriptionPlugin(BasePlugin):
+    def get_schema(
+        self,
+        instance: Instance,
+        ctx: Context,
+        schema: Optional[JSONSchema] = None,
+    ) -> Optional[JSONSchema]:
+        if schema and is_dataclass(instance.type) and instance.type.__doc__:
+            schema.description = cleandoc(instance.type.__doc__)
+        return None
diff --git a/mashumaro/jsonschema/schema.py b/mashumaro/jsonschema/schema.py
index 56e0f49..dfe6bde 100644
--- a/mashumaro/jsonschema/schema.py
+++ b/mashumaro/jsonschema/schema.py
@@ -105,6 +105,9 @@ class Instance:
     __owner_builder: Optional[CodeBuilder] = None
     __self_builder: Optional[CodeBuilder] = None
 
+    # Original type despite custom serialization. To be revised.
+    _original_type: Type = field(init=False)
+
     origin_type: Type = field(init=False)
     annotations: list[Annotation] = field(init=False, default_factory=list)
 
@@ -150,6 +153,7 @@ def derive(self, **changes: Any) -> "Instance":
         return new_instance
 
     def __post_init__(self) -> None:
+        self._original_type = self.type
         self.update_type(self.type)
         if is_annotated(self.type):
             self.annotations = getattr(self.type, "__metadata__", [])
@@ -260,12 +264,22 @@ class EmptyJSONSchema(JSONSchema):
 def get_schema(
     instance: Instance, ctx: Context, with_dialect_uri: bool = False
 ) -> JSONSchema:
+    schema = None
     for schema_creator in Registry.iter():
         schema = schema_creator(instance, ctx)
         if schema is not None:
             if with_dialect_uri:
                 schema.schema = ctx.dialect.uri
-            return schema
+            break
+    for plugin in ctx.plugins:
+        try:
+            new_schema = plugin.get_schema(instance, ctx, schema)
+            if new_schema:
+                schema = new_schema
+        except NotImplementedError:
+            continue
+    if schema:
+        return schema
     raise NotImplementedError(
         f'Type {type_name(instance.type)} of field "{instance.name}" '
         f"in {type_name(instance.owner_class)} isn't supported"


Test patch:
diff --git a/tests/test_jsonschema/test_jsonschema_plugins.py b/tests/test_jsonschema/test_jsonschema_plugins.py
new file mode 100644
index 0000000..2723b46
--- /dev/null
+++ b/tests/test_jsonschema/test_jsonschema_plugins.py
@@ -0,0 +1,127 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import pytest
+
+from mashumaro.jsonschema.builder import JSONSchemaBuilder, build_json_schema
+from mashumaro.jsonschema.models import (
+    Context,
+    JSONSchema,
+    JSONSchemaInstanceType,
+)
+from mashumaro.jsonschema.plugins import BasePlugin, DocstringDescriptionPlugin
+from mashumaro.jsonschema.schema import Instance
+
+
+class ThirdPartyType:
+    pass
+
+
+@dataclass
+class DataClassWithDocstring:
+    """Here is the docstring"""
+
+    x: int
+
+
+@dataclass
+class DataClassWithoutDocstring:
+    x: int
+
+
+def test_basic_plugin():
+    assert build_json_schema(int, plugins=[BasePlugin()]) == JSONSchema(
+        type=JSONSchemaInstanceType.INTEGER
+    )
+
+
+def test_plugin_with_not_implemented_error():
+    class NotImplementedErrorPlugin(BasePlugin):
+        def get_schema(
+            self,
+            instance: Instance,
+            ctx: Context,
+            schema: Optional[JSONSchema] = None,
+        ) -> Optional[JSONSchema]:
+            raise NotImplementedError
+
+    assert build_json_schema(
+        int, plugins=[NotImplementedErrorPlugin()]
+    ) == JSONSchema(type=JSONSchemaInstanceType.INTEGER)
+    assert JSONSchemaBuilder(plugins=[NotImplementedErrorPlugin()]).build(
+        int
+    ) == JSONSchema(type=JSONSchemaInstanceType.INTEGER)
+
+
+@pytest.mark.parametrize(
+    ("obj", "docstring"),
+    (
+        (DataClassWithDocstring, "Here is the docstring"),
+        (DataClassWithoutDocstring, "DataClassWithoutDocstring(x: int)"),
+        (int, None),
+    ),
+)
+def test_docstring_description_plugin(obj, docstring):
+    assert build_json_schema(obj).description is None
+    assert JSONSchemaBuilder().build(obj).description is None
+
+    assert (
+        build_json_schema(
+            obj, plugins=[DocstringDescriptionPlugin()]
+        ).description
+        == docstring
+    )
+    assert (
+        JSONSchemaBuilder(plugins=[DocstringDescriptionPlugin()])
+        .build(obj)
+        .description
+        == docstring
+    )
+
+
+def test_third_party_type_plugin():
+    third_party_json_schema = JSONSchema()
+
+    class ThirdPartyTypePlugin(BasePlugin):
+        def get_schema(
+            self,
+            instance: Instance,
+            ctx: Context,
+            schema: Optional[JSONSchema] = None,
+        ) -> Optional[JSONSchema]:
+            try:
+                if issubclass(instance.type, ThirdPartyType):
+                    return third_party_json_schema
+            except TypeError:
+                pass
+
+    assert (
+        build_json_schema(ThirdPartyType, plugins=[ThirdPartyTypePlugin()])
+        is third_party_json_schema
+    )
+    assert (
+        JSONSchemaBuilder(plugins=[ThirdPartyTypePlugin()]).build(
+            ThirdPartyType
+        )
+        is third_party_json_schema
+    )
+    assert (
+        JSONSchemaBuilder(plugins=[ThirdPartyTypePlugin()])
+        .build(list[ThirdPartyType])
+        .items
+        is third_party_json_schema
+    )
+    assert (
+        build_json_schema(
+            list[ThirdPartyType], plugins=[ThirdPartyTypePlugin()]
+        ).items
+        is third_party_json_schema
+    )
+    with pytest.raises(NotImplementedError):
+        build_json_schema(ThirdPartyType)
+    with pytest.raises(NotImplementedError):
+        JSONSchemaBuilder().build(ThirdPartyType)
+    with pytest.raises(NotImplementedError):
+        build_json_schema(list[ThirdPartyType])
+    with pytest.raises(NotImplementedError):
+        JSONSchemaBuilder().build(list[ThirdPartyType])