JsonSchemaTransformer: Add back support for the simplify_nullable_unions kwarg

Author: conradlee

pydantic_ai_slim/pydantic_ai/_json_schema.py

@@ -25,13 +25,15 @@ def __init__(
         *,
         strict: bool | None = None,
         prefer_inlined_defs: bool = False,
+        simplify_nullable_unions: bool = False,  # TODO (v2): Remove this, no longer used
     ):
         self.schema = schema
 
         self.strict = strict
         self.is_strict_compatible = True  # Can be set to False by subclasses to set `strict` on `ToolDefinition` when set not set by user explicitly
 
         self.prefer_inlined_defs = prefer_inlined_defs
+        self.simplify_nullable_unions = simplify_nullable_unions  # TODO (v2): Remove this, no longer used
 
         self.defs: dict[str, JsonSchema] = self.schema.get('$defs', {})
         self.refs_stack: list[str] = []
@@ -144,11 +146,39 @@ def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf'
 
         handled = [self._handle(member) for member in members]
 
+        # TODO (v2): Remove this feature, no longer used
+        if self.simplify_nullable_unions:
+            handled = self._simplify_nullable_union(handled)
+
+        if len(handled) == 1:
+            # In this case, no need to retain the union
+            return handled[0] | schema
+
         # If we have keys besides the union kind (such as title or discriminator), keep them without modifications
         schema = schema.copy()
         schema[union_kind] = handled
         return schema
 
+    @staticmethod
+    def _simplify_nullable_union(cases: list[JsonSchema]) -> list[JsonSchema]:
+        # TODO (v2): Remove this method, no longer used
+        if len(cases) == 2 and {'type': 'null'} in cases:
+            # Find the non-null schema
+            non_null_schema = next(
+                (item for item in cases if item != {'type': 'null'}),
+                None,
+            )
+            if non_null_schema:
+                # Create a new schema based on the non-null part, mark as nullable
+                new_schema = deepcopy(non_null_schema)
+                new_schema['nullable'] = True
+                return [new_schema]
+            else:  # pragma: no cover
+                # they are both null, so just return one of them
+                return [cases[0]]
+
+        return cases
+
 
 class InlineDefsJsonSchemaTransformer(JsonSchemaTransformer):
     """Transforms the JSON Schema to inline $defs."""

tests/test_utils.py

@@ -6,11 +6,13 @@
 import os
 from collections.abc import AsyncIterator
 from importlib.metadata import distributions
+from typing import Any
 
 import pytest
 from inline_snapshot import snapshot
 
 from pydantic_ai import UserError
+from pydantic_ai._json_schema import JsonSchemaTransformer
 from pydantic_ai._utils import (
     UNSET,
     PeekableAsyncStream,
@@ -541,3 +543,47 @@ def test_validate_empty_kwargs_preserves_order():
     assert '`first`' in error_msg
     assert '`second`' in error_msg
     assert '`third`' in error_msg
+
+
+def test_simplify_nullable_unions():
+    """Test the simplify_nullable_unions feature (deprecated, to be removed in v2)."""
+
+    # Create a concrete subclass for testing
+    class TestTransformer(JsonSchemaTransformer):
+        def transform(self, schema: dict[str, Any]) -> dict[str, Any]:
+            return schema
+
+    # Test with simplify_nullable_unions=True
+    schema_with_null = {
+        'anyOf': [
+            {'type': 'string'},
+            {'type': 'null'},
+        ]
+    }
+    transformer = TestTransformer(schema_with_null, simplify_nullable_unions=True)
+    result = transformer.walk()
+
+    # Should collapse to a single nullable string
+    assert result == {'type': 'string', 'nullable': True}
+
+    # Test with simplify_nullable_unions=False (default)
+    transformer2 = TestTransformer(schema_with_null, simplify_nullable_unions=False)
+    result2 = transformer2.walk()
+
+    # Should keep the anyOf structure
+    assert 'anyOf' in result2
+    assert len(result2['anyOf']) == 2
+
+    # Test that non-nullable unions are unaffected
+    schema_no_null = {
+        'anyOf': [
+            {'type': 'string'},
+            {'type': 'number'},
+        ]
+    }
+    transformer3 = TestTransformer(schema_no_null, simplify_nullable_unions=True)
+    result3 = transformer3.walk()
+
+    # Should keep anyOf since it's not nullable
+    assert 'anyOf' in result3
+    assert len(result3['anyOf']) == 2