Expose public sort method for JSON schema generation (#10595)

sydney-runkle · Viicos · web-flow · commit 07e2b729af2d · 2024-10-10T11:39:08.000-04:00
Co-authored-by: Victorien &lt;65306057+Viicos@users.noreply.github.com&gt;
diff --git a/docs/concepts/json_schema.md b/docs/concepts/json_schema.md
@@ -1414,6 +1414,68 @@ print(validation_schema)
 """
 ```
 
+### JSON schema sorting
+
+By default, Pydantic recursively sorts JSON schemas by alphabetically sorting keys. Notably, Pydantic skips sorting the values of the `properties` key,
+to preserve the order of the fields as they were defined in the model.
+
+If you would like to customize this behavior, you can override the `sort` method in your custom `GenerateJsonSchema` subclass. The below example
+uses a no-op `sort` method to disable sorting entirely, which is reflected in the preserved order of the model fields and `json_schema_extra` keys:
+
+```py
+import json
+from typing import Optional
+
+from pydantic import BaseModel, Field
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
+
+
+class MyGenerateJsonSchema(GenerateJsonSchema):
+    def sort(
+        self, value: JsonSchemaValue, parent_key: Optional[str] = None
+    ) -> JsonSchemaValue:
+        """No-op, we don't want to sort schema values at all."""
+        return value
+
+
+class Bar(BaseModel):
+    c: str
+    b: str
+    a: str = Field(json_schema_extra={'c': 'hi', 'b': 'hello', 'a': 'world'})
+
+
+json_schema = Bar.model_json_schema(schema_generator=MyGenerateJsonSchema)
+print(json.dumps(json_schema, indent=2))
+"""
+{
+  "type": "object",
+  "properties": {
+    "c": {
+      "type": "string",
+      "title": "C"
+    },
+    "b": {
+      "type": "string",
+      "title": "B"
+    },
+    "a": {
+      "type": "string",
+      "c": "hi",
+      "b": "hello",
+      "a": "world",
+      "title": "A"
+    }
+  },
+  "required": [
+    "c",
+    "b",
+    "a"
+  ],
+  "title": "Bar"
+}
+"""
+```
+
 ## Customizing the `$ref`s in JSON Schema
 
 The format of `$ref`s can be altered by calling [`model_json_schema()`][pydantic.main.BaseModel.model_json_schema]
diff --git a/pydantic/json_schema.py b/pydantic/json_schema.py
@@ -392,7 +392,7 @@ def generate_definitions(
         json_schema = {'$defs': self.definitions}
         json_schema = definitions_remapping.remap_json_schema(json_schema)
         self._used = True
-        return json_schemas_map, _sort_json_schema(json_schema['$defs'])  # type: ignore
+        return json_schemas_map, self.sort(json_schema['$defs'])  # type: ignore
 
     def generate(self, schema: CoreSchema, mode: JsonSchemaMode = 'validation') -> JsonSchemaValue:
         """Generates a JSON schema for a specified schema in a specified mode.
@@ -441,7 +441,7 @@ def generate(self, schema: CoreSchema, mode: JsonSchemaMode = 'validation') -> J
         # json_schema['$schema'] = self.schema_dialect
 
         self._used = True
-        return _sort_json_schema(json_schema)
+        return self.sort(json_schema)
 
     def generate_inner(self, schema: CoreSchemaOrField) -> JsonSchemaValue:  # noqa: C901
         """Generates a JSON schema for a given core schema.
@@ -556,6 +556,38 @@ def new_handler_func(
             json_schema = populate_defs(schema, json_schema)
         return json_schema
 
+    def sort(self, value: JsonSchemaValue, parent_key: str | None = None) -> JsonSchemaValue:
+        """Override this method to customize the sorting of the JSON schema (e.g., don't sort at all, sort all keys unconditionally, etc.)
+
+        By default, alphabetically sort the keys in the JSON schema, skipping the 'properties' and 'default' keys to preserve field definition order.
+        This sort is recursive, so it will sort all nested dictionaries as well.
+        """
+        sorted_dict: dict[str, JsonSchemaValue] = {}
+        keys = value.keys()
+        if parent_key not in ('properties', 'default'):
+            keys = sorted(keys)
+        for key in keys:
+            sorted_dict[key] = self._sort_recursive(value[key], parent_key=key)
+        return sorted_dict
+
+    def _sort_recursive(self, value: Any, parent_key: str | None = None) -> Any:
+        """Recursively sort a JSON schema value."""
+        if isinstance(value, dict):
+            sorted_dict: dict[str, JsonSchemaValue] = {}
+            keys = value.keys()
+            if parent_key not in ('properties', 'default'):
+                keys = sorted(keys)
+            for key in keys:
+                sorted_dict[key] = self._sort_recursive(value[key], parent_key=key)
+            return sorted_dict
+        elif isinstance(value, list):
+            sorted_list: list[JsonSchemaValue] = []
+            for item in value:
+                sorted_list.append(self._sort_recursive(item, parent_key))
+            return sorted_list
+        else:
+            return value
+
     # ### Schema generation methods
     def any_schema(self, schema: core_schema.AnySchema) -> JsonSchemaValue:
         """Generates a JSON schema that matches any value.
@@ -2367,24 +2399,6 @@ def _make_json_hashable(value: JsonValue) -> _HashableJsonValue:
         return value
 
 
-def _sort_json_schema(value: JsonSchemaValue, parent_key: str | None = None) -> JsonSchemaValue:
-    if isinstance(value, dict):
-        sorted_dict: dict[str, JsonSchemaValue] = {}
-        keys = value.keys()
-        if (parent_key != 'properties') and (parent_key != 'default'):
-            keys = sorted(keys)
-        for key in keys:
-            sorted_dict[key] = _sort_json_schema(value[key], parent_key=key)
-        return sorted_dict
-    elif isinstance(value, list):
-        sorted_list: list[JsonSchemaValue] = []
-        for item in value:  # type: ignore
-            sorted_list.append(_sort_json_schema(item, parent_key))
-        return sorted_list  # type: ignore
-    else:
-        return value
-
-
 @dataclasses.dataclass(**_internal_dataclass.slots_true)
 class WithJsonSchema:
     """Usage docs: https://docs.pydantic.dev/2.10/concepts/json_schema/#withjsonschema-annotation
