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