Operate on a deepcopy of $defs in JsonSchemaTransformer instead of the original schema (#3758)

Author: thomasst

pydantic_ai_slim/pydantic_ai/_json_schema.py

@@ -45,7 +45,7 @@ def __init__(
  self.prefer_inlined_defs = prefer_inlined_defs
  self.simplify_nullable_unions = simplify_nullable_unions
 
- self.defs: dict[str, JsonSchema] = self.schema.get('$defs', {})
+ self.defs: dict[str, JsonSchema] = deepcopy(self.schema.get('$defs', {}))
  self.refs_stack: list[str] = []
  self.recursive_refs = set[str]()
 

tests/test_json_schema.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations as _annotations
 
+from copy import deepcopy
 from typing import Any
 
 from pydantic_ai._json_schema import JsonSchemaTransformer
@@ -49,3 +50,40 @@ def transform(self, schema: dict[str, Any]) -> dict[str, Any]:
  # Should keep anyOf since it's not nullable
  assert 'anyOf' in result3
  assert len(result3['anyOf']) == 2
+
+
+def test_schema_defs_not_modified():
+ """Test that the original schema $defs are not modified during transformation."""
+
+ # Create a concrete subclass for testing
+ class TestTransformer(JsonSchemaTransformer):
+ def transform(self, schema: dict[str, Any]) -> dict[str, Any]:
+ return schema
+
+ # Create a schema with $defs that should not be modified
+ original_schema = {
+ 'type': 'object',
+ 'properties': {'value': {'$ref': '#/$defs/TestUnion'}},
+ '$defs': {
+ 'TestUnion': {
+ 'anyOf': [
+ {'type': 'string'},
+ {'type': 'number'},
+ ],
+ 'title': 'TestUnion',
+ }
+ },
+ }
+
+ # Keep a deepcopy to compare against later
+ original_schema_copy = deepcopy(original_schema)
+
+ # Transform the schema
+ transformer = TestTransformer(original_schema)
+ result = transformer.walk()
+
+ # Verify the original schema was not modified
+ assert original_schema == original_schema_copy
+
+ # Verify the result is correct
+ assert result == original_schema_copy