Merge pull request #265 from Fatal1ty/json-schema-plugins

Add a plugin system to json schema

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.

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:

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] = ()

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

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"