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())
+```
+
+
+Click to show the result
+
+```json
+{
+ "type": "object",
+ "title": "MyClass",
+ "description": "My class",
+ "properties": {
+ "x": {
+ "type": "integer"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "x"
+ ]
+}
+```
+
+
+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())
+```
+
+
+Click to show the result
+
+```json
+{
+ "type": "object",
+ "title": "MyDataClass",
+ "properties": {
+ "y": {
+ "type": "object",
+ "title": "MyPydanticClass",
+ "properties": {
+ "x": {
+ "type": "integer",
+ "title": "X"
+ }
+ },
+ "required": [
+ "x"
+ ]
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "y"
+ ]
+}
+```
+
+
+
### 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"
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])