Simplify, report all ref names

pydantic · Oct 22, 2024 · e75dbe7 · e75dbe7
1 parent 8a3ade1
commit e75dbe7
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 13 deletions.
diff --git a/python/pydantic_core/__init__.py b/python/pydantic_core/__init__.py
@@ -146,6 +146,6 @@ class MultiHostHost(_TypedDict):
 class GatherResult(_TypedDict):
     """Internal result of gathering schemas for cleaning."""
 
-    inlinable_def_refs: dict[str, DefinitionReferenceSchema]
+    inlinable_def_refs: dict[str, DefinitionReferenceSchema | None]
     recursive_refs: set[str]
     schemas_with_meta_keys: dict[str, list[CoreSchema]] | None
diff --git a/src/schema_traverse.rs b/src/schema_traverse.rs
@@ -193,15 +193,8 @@ pub fn gather_schemas_for_cleaning<'py>(
     };
     gather_schema(schema.downcast_exact()?, &mut ctx)?;
 
-    let inlinable_def_refs = PyDict::new_bound(py);
-    for (ref_str, def_ref_candidate) in ctx.inline_def_ref_candidates.iter() {
-        if !def_ref_candidate.is_none() {
-            inlinable_def_refs.set_item(ref_str, def_ref_candidate)?;
-        }
-    }
-
     let res = PyDict::new_bound(py);
-    res.set_item(intern!(py, "inlinable_def_refs"), inlinable_def_refs)?;
+    res.set_item(intern!(py, "inlinable_def_refs"), ctx.inline_def_ref_candidates)?;
     res.set_item(intern!(py, "recursive_refs"), ctx.recursive_def_refs)?;
     res.set_item(intern!(py, "schemas_with_meta_keys"), ctx.meta_with_keys.map(|v| v.0))?;
     Ok(res)

diff --git a/tests/test_gather_schemas_for_cleaning.py b/tests/test_gather_schemas_for_cleaning.py
@@ -42,7 +42,7 @@ class Model:
     definitions = {'ref1': core_schema.str_schema(ref='ref1'), 'ref2': core_schema.bytes_schema(ref='ref2')}
 
     res = gather_schemas_for_cleaning(schema, definitions, find_meta_with_keys=None)
-    assert res['inlinable_def_refs'] == {'ref2': ref2} and res['inlinable_def_refs']['ref2'] is ref2
+    assert res['inlinable_def_refs'] == {'ref1': None, 'ref2': ref2} and res['inlinable_def_refs']['ref2'] is ref2
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] is None
 
@@ -61,7 +61,7 @@ def test_ref_in_serialization_schema():
 def test_recursive_ref_schema():
     ref1 = core_schema.definition_reference_schema('ref1')
     res = gather_schemas_for_cleaning(ref1, definitions={'ref1': ref1}, find_meta_with_keys=None)
-    assert res['inlinable_def_refs'] == {}
+    assert res['inlinable_def_refs'] == {'ref1': None}
     assert res['recursive_refs'] == {'ref1'}
     assert res['schemas_with_meta_keys'] is None
 
@@ -80,7 +80,7 @@ def test_deep_recursive_ref_schema():
         },
         find_meta_with_keys=None,
     )
-    assert res['inlinable_def_refs'] == {}
+    assert res['inlinable_def_refs'] == {'ref1': None, 'ref2': None, 'ref3': None}
     assert res['recursive_refs'] == {'ref1', 'ref2', 'ref3'}
     assert res['schemas_with_meta_keys'] is None
 
@@ -129,6 +129,6 @@ def test_no_duplicate_ref_instances_gathered():
     definitions = {'ref1': schema1, 'ref2': schema2}
 
     res = gather_schemas_for_cleaning(schema3, definitions=definitions, find_meta_with_keys=None)
-    assert res['inlinable_def_refs'] == {}
+    assert res['inlinable_def_refs'] == {'ref1': None, 'ref2': None}
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] is None