diff --git a/dandischema/utils.py b/dandischema/utils.py index 33ebfe6..c73167c 100644 --- a/dandischema/utils.py +++ b/dandischema/utils.py @@ -3,8 +3,8 @@ import re from typing import Any, Iterator, List, Union, get_args, get_origin -from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue -from pydantic_core import core_schema +from pydantic.json_schema import GenerateJsonSchema, JsonSchemaMode, JsonSchemaValue +from pydantic_core import CoreSchema, core_schema TITLE_CASE_LOWER = { "a", @@ -69,6 +69,16 @@ class TransitionalGenerateJsonSchema(GenerateJsonSchema): schema generation process are the same as the behavior in Pydantic V1. """ + def generate( + self, schema: CoreSchema, mode: JsonSchemaMode = "validation" + ) -> JsonSchemaValue: + json_schema = super().generate(schema, mode) + + # Set the `$schema` key with the schema dialect + json_schema["$schema"] = self.schema_dialect + + return json_schema + def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue: # Override the default behavior for handling a schema that allows null values # With this override, `Optional` fields will not be indicated