From 5f89fcd1e9a2130cce938ffb9c4a614612821853 Mon Sep 17 00:00:00 2001 From: Alexander Tikhonov Date: Thu, 21 Nov 2024 20:13:32 +0300 Subject: [PATCH 1/7] Add a plugin system to json schema --- mashumaro/jsonschema/builder.py | 6 +++++- mashumaro/jsonschema/models.py | 9 ++++++++- mashumaro/jsonschema/plugins.py | 26 ++++++++++++++++++++++++++ mashumaro/jsonschema/schema.py | 12 +++++++++++- 4 files changed, 50 insertions(+), 3 deletions(-) create mode 100644 mashumaro/jsonschema/plugins.py diff --git a/mashumaro/jsonschema/builder.py b/mashumaro/jsonschema/builder.py index 23a31cb..4f3cf22 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,15 +23,17 @@ 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() + context = Context(plugins=plugins) else: context = Context( dialect=context.dialect, definitions=context.definitions, all_refs=context.all_refs, ref_prefix=context.ref_prefix, + plugins=plugins, ) if dialect is not None: context.dialect = dialect diff --git a/mashumaro/jsonschema/models.py b/mashumaro/jsonschema/models.py index 3d0f785..d38d3c3 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 Self, TypeAlias, TYPE_CHECKING 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..d19fab1 --- /dev/null +++ b/mashumaro/jsonschema/plugins.py @@ -0,0 +1,26 @@ +from dataclasses import is_dataclass +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, + ): + pass + + +class DocstringDescriptionPlugin(BasePlugin): + def get_schema( + self, + instance: Instance, + ctx: Context, + schema: Optional[JSONSchema] = None, + ): + if schema and is_dataclass(instance.type) and instance.type.__doc__: + schema.description = instance.type.__doc__ diff --git a/mashumaro/jsonschema/schema.py b/mashumaro/jsonschema/schema.py index 56e0f49..1145216 100644 --- a/mashumaro/jsonschema/schema.py +++ b/mashumaro/jsonschema/schema.py @@ -260,12 +260,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" From f6de71be03a3042509d2d7e15f0d44e6138f4671 Mon Sep 17 00:00:00 2001 From: Alexander Tikhonov Date: Thu, 21 Nov 2024 20:21:06 +0300 Subject: [PATCH 2/7] Fixes --- mashumaro/jsonschema/builder.py | 8 ++++++-- mashumaro/jsonschema/models.py | 2 +- mashumaro/jsonschema/plugins.py | 5 +++-- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/mashumaro/jsonschema/builder.py b/mashumaro/jsonschema/builder.py index 4f3cf22..28a23ab 100644 --- a/mashumaro/jsonschema/builder.py +++ b/mashumaro/jsonschema/builder.py @@ -26,14 +26,14 @@ def build_json_schema( plugins: Sequence[BasePlugin] = (), ) -> JSONSchema: if context is None: - context = Context(plugins=plugins) + context = Context() else: context = Context( dialect=context.dialect, definitions=context.definitions, all_refs=context.all_refs, ref_prefix=context.ref_prefix, - plugins=plugins, + plugins=context.plugins, ) if dialect is not None: context.dialect = dialect @@ -45,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: @@ -68,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 @@ -77,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 d38d3c3..b194ccc 100644 --- a/mashumaro/jsonschema/models.py +++ b/mashumaro/jsonschema/models.py @@ -5,7 +5,7 @@ from enum import Enum from typing import Any, Dict, List, Optional, Union -from typing_extensions import Self, TypeAlias, TYPE_CHECKING +from typing_extensions import TYPE_CHECKING, Self, TypeAlias from mashumaro.config import BaseConfig from mashumaro.helper import pass_through diff --git a/mashumaro/jsonschema/plugins.py b/mashumaro/jsonschema/plugins.py index d19fab1..0974d07 100644 --- a/mashumaro/jsonschema/plugins.py +++ b/mashumaro/jsonschema/plugins.py @@ -11,7 +11,7 @@ def get_schema( instance: Instance, ctx: Context, schema: Optional[JSONSchema] = None, - ): + ) -> Optional[JSONSchema]: pass @@ -21,6 +21,7 @@ def get_schema( instance: Instance, ctx: Context, schema: Optional[JSONSchema] = None, - ): + ) -> Optional[JSONSchema]: if schema and is_dataclass(instance.type) and instance.type.__doc__: schema.description = instance.type.__doc__ + return None From 931e63205c689c0f3f23c600d7b2579e5d2d20e4 Mon Sep 17 00:00:00 2001 From: Alexander Tikhonov Date: Thu, 21 Nov 2024 23:38:17 +0300 Subject: [PATCH 3/7] Update README --- README.md | 124 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) 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. From 24403f790d0f295d39ad03aeb697ac5d479e769d Mon Sep 17 00:00:00 2001 From: Alexander Tikhonov Date: Sat, 23 Nov 2024 14:19:02 +0300 Subject: [PATCH 4/7] Add tests for JSON Schema plugins --- .../test_jsonschema_plugins.py | 99 +++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 tests/test_jsonschema/test_jsonschema_plugins.py diff --git a/tests/test_jsonschema/test_jsonschema_plugins.py b/tests/test_jsonschema/test_jsonschema_plugins.py new file mode 100644 index 0000000..2a58bfb --- /dev/null +++ b/tests/test_jsonschema/test_jsonschema_plugins.py @@ -0,0 +1,99 @@ +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 +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 + + +@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]) From fcfe2d398b6462b9834fa4af3d80ff109b33152b Mon Sep 17 00:00:00 2001 From: Alexander Tikhonov Date: Sat, 23 Nov 2024 14:30:27 +0300 Subject: [PATCH 5/7] Add more tests --- .../test_jsonschema_plugins.py | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/test_jsonschema/test_jsonschema_plugins.py b/tests/test_jsonschema/test_jsonschema_plugins.py index 2a58bfb..2723b46 100644 --- a/tests/test_jsonschema/test_jsonschema_plugins.py +++ b/tests/test_jsonschema/test_jsonschema_plugins.py @@ -4,7 +4,11 @@ import pytest from mashumaro.jsonschema.builder import JSONSchemaBuilder, build_json_schema -from mashumaro.jsonschema.models import Context, JSONSchema +from mashumaro.jsonschema.models import ( + Context, + JSONSchema, + JSONSchemaInstanceType, +) from mashumaro.jsonschema.plugins import BasePlugin, DocstringDescriptionPlugin from mashumaro.jsonschema.schema import Instance @@ -25,6 +29,30 @@ 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"), ( From 98c1a277e28beed12e3291ec48cd95bba8a37b19 Mon Sep 17 00:00:00 2001 From: Alexander Tikhonov Date: Sat, 23 Nov 2024 15:48:50 +0300 Subject: [PATCH 6/7] Add a temporary workaround to get original type despite custom serialization --- mashumaro/jsonschema/schema.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mashumaro/jsonschema/schema.py b/mashumaro/jsonschema/schema.py index 1145216..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__", []) From c226c0f03aefebd21d7b99762e5f3e6d790ac747 Mon Sep 17 00:00:00 2001 From: Alexander Tikhonov Date: Sat, 23 Nov 2024 15:54:14 +0300 Subject: [PATCH 7/7] Use cleandoc in DocstringDescriptionPlugin --- mashumaro/jsonschema/plugins.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mashumaro/jsonschema/plugins.py b/mashumaro/jsonschema/plugins.py index 0974d07..3396c9c 100644 --- a/mashumaro/jsonschema/plugins.py +++ b/mashumaro/jsonschema/plugins.py @@ -1,4 +1,5 @@ from dataclasses import is_dataclass +from inspect import cleandoc from typing import Optional from mashumaro.jsonschema.models import Context, JSONSchema @@ -23,5 +24,5 @@ def get_schema( schema: Optional[JSONSchema] = None, ) -> Optional[JSONSchema]: if schema and is_dataclass(instance.type) and instance.type.__doc__: - schema.description = instance.type.__doc__ + schema.description = cleandoc(instance.type.__doc__) return None