Move discriminator const away. Use generic finding

Author: MarkusSintonen

python/pydantic_core/__init__.py

@@ -146,4 +146,4 @@ class GatherResult(_TypedDict):
 
     definition_refs: dict[str, list[DefinitionReferenceSchema]]
     recursive_refs: set[str]
-    deferred_discriminators: list[tuple[CoreSchema, _Any]]
+    schemas_with_meta_keys: dict[str, list[CoreSchema]] | None

python/pydantic_core/_pydantic_core.pyi

@@ -1166,6 +1166,10 @@ def validate_core_schema(schema: CoreSchema, *, strict: bool | None = None) -> C
     using pydantic-core directly.
     """
 
-def gather_schemas_for_cleaning(schema: CoreSchema, definitions: dict[str, CoreSchema]) -> GatherResult:
+def gather_schemas_for_cleaning(
+    schema: CoreSchema,
+    definitions: dict[str, CoreSchema],
+    find_meta_with_keys: set[str] | None,
+) -> GatherResult:
     """Used internally for schema cleaning when schemas are generated.
     Gathers information from the schema tree for the cleaning."""

src/schema_traverse.rs

@@ -5,8 +5,6 @@ use pyo3::types::{PyDict, PyList, PySet, PyString, PyTuple};
 use pyo3::{intern, Bound, PyResult};
 use std::collections::HashSet;
 
-const CORE_SCHEMA_METADATA_DISCRIMINATOR_PLACEHOLDER_KEY: &str = "pydantic.internal.union_discriminator";
-
 macro_rules! get {
     ($dict: expr, $key: expr) => {
         $dict.get_item(intern!($dict.py(), $key))?
@@ -52,7 +50,7 @@ fn gather_definition_ref(schema_ref_dict: &Bound<'_, PyDict>, ctx: &mut GatherCt
             defaultdict_list_append!(&ctx.def_refs, schema_ref_pystr, schema_ref_dict);
 
             // TODO should py_err! when not found. That error can be used to detect the missing defs in cleaning side
-            if let Some(definition) = ctx.definitions_dict.get_item(schema_ref_pystr)? {
+            if let Some(definition) = ctx.definitions.get_item(schema_ref_pystr)? {
                 ctx.recursively_seen_refs.insert(schema_ref_str.to_string());
 
                 gather_schema(definition.downcast_exact::<PyDict>()?, ctx)?;
@@ -75,11 +73,13 @@ fn gather_definition_ref(schema_ref_dict: &Bound<'_, PyDict>, ctx: &mut GatherCt
 }
 
 fn gather_meta(schema: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()> {
-    if let Some(meta) = get!(schema, "metadata") {
-        let meta_dict = meta.downcast_exact::<PyDict>()?;
-        if let Some(discriminator) = get!(meta_dict, CORE_SCHEMA_METADATA_DISCRIMINATOR_PLACEHOLDER_KEY) {
-            let schema_discriminator = PyTuple::new_bound(schema.py(), vec![schema.as_any(), &discriminator]);
-            ctx.discriminators.append(schema_discriminator)?;
+    if let Some((res, find_keys)) = &ctx.meta_with_keys {
+        if let Some(meta) = get!(schema, "metadata") {
+            for (k, _) in meta.downcast_exact::<PyDict>()?.iter() {
+                if find_keys.contains(&k)? {
+                    defaultdict_list_append!(res, &k, schema);
+                }
+            }
         }
     }
     Ok(())
@@ -152,32 +152,38 @@ fn gather_schema(schema: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()
     }
 }
 
-pub struct GatherCtx<'a, 'py> {
-    pub definitions_dict: &'a Bound<'py, PyDict>,
-    pub def_refs: Bound<'py, PyDict>,
-    pub recursive_def_refs: Bound<'py, PySet>,
-    pub discriminators: Bound<'py, PyList>,
+struct GatherCtx<'a, 'py> {
+    definitions: &'a Bound<'py, PyDict>,
+    meta_with_keys: Option<(Bound<'py, PyDict>, &'a Bound<'py, PySet>)>,
+    def_refs: Bound<'py, PyDict>,
+    recursive_def_refs: Bound<'py, PySet>,
     recursively_seen_refs: HashSet<String>,
 }
 
-#[pyfunction(signature = (schema, definitions))]
+#[pyfunction(signature = (schema, definitions, find_meta_with_keys))]
 pub fn gather_schemas_for_cleaning<'py>(
     schema: &Bound<'py, PyAny>,
     definitions: &Bound<'py, PyAny>,
+    find_meta_with_keys: &Bound<'py, PyAny>,
 ) -> PyResult<Bound<'py, PyDict>> {
     let py = schema.py();
+    let meta_with_keys = if find_meta_with_keys.is_none() {
+        None
+    } else {
+        Some((PyDict::new_bound(py), find_meta_with_keys.downcast_exact::<PySet>()?))
+    };
     let mut ctx = GatherCtx {
-        definitions_dict: definitions.downcast_exact()?,
+        definitions: definitions.downcast_exact()?,
+        meta_with_keys,
         def_refs: PyDict::new_bound(py),
         recursive_def_refs: PySet::empty_bound(py)?,
-        discriminators: PyList::empty_bound(py),
         recursively_seen_refs: HashSet::new(),
     };
-    gather_schema(schema.downcast_exact::<PyDict>()?, &mut ctx)?;
+    gather_schema(schema.downcast_exact()?, &mut ctx)?;
 
     let res = PyDict::new_bound(py);
     res.set_item(intern!(py, "definition_refs"), ctx.def_refs)?;
     res.set_item(intern!(py, "recursive_refs"), ctx.recursive_def_refs)?;
-    res.set_item(intern!(py, "deferred_discriminators"), ctx.discriminators)?;
+    res.set_item(intern!(py, "schemas_with_meta_keys"), ctx.meta_with_keys.map(|v| v.0))?;
     Ok(res)
 }

tests/test_gather_schemas_for_cleaning.py

@@ -5,20 +5,20 @@ def test_no_refs():
     p1 = core_schema.arguments_parameter('a', core_schema.int_schema())
     p2 = core_schema.arguments_parameter('b', core_schema.int_schema())
     schema = core_schema.arguments_schema([p1, p2])
-    res = gather_schemas_for_cleaning(schema, definitions={})
+    res = gather_schemas_for_cleaning(schema, definitions={}, find_meta_with_keys=None)
     assert res['definition_refs'] == {}
     assert res['recursive_refs'] == set()
-    assert res['deferred_discriminators'] == []
+    assert res['schemas_with_meta_keys'] is None
 
 
 def test_simple_ref_schema():
     schema = core_schema.definition_reference_schema('ref1')
     definitions = {'ref1': core_schema.int_schema(ref='ref1')}
 
-    res = gather_schemas_for_cleaning(schema, definitions)
+    res = gather_schemas_for_cleaning(schema, definitions, find_meta_with_keys=None)
     assert res['definition_refs'] == {'ref1': [schema]} and res['definition_refs']['ref1'][0] is schema
     assert res['recursive_refs'] == set()
-    assert res['deferred_discriminators'] == []
+    assert res['schemas_with_meta_keys'] is None
 
 
 def test_deep_ref_schema():
@@ -39,31 +39,31 @@ class Model:
     )
     definitions = {'ref1': core_schema.str_schema(ref='ref1'), 'ref2': core_schema.bytes_schema(ref='ref2')}
 
-    res = gather_schemas_for_cleaning(schema, definitions)
+    res = gather_schemas_for_cleaning(schema, definitions, find_meta_with_keys=None)
     assert res['definition_refs'] == {'ref1': [ref11, ref12], 'ref2': [ref2]}
     assert res['definition_refs']['ref1'][0] is ref11 and res['definition_refs']['ref1'][1] is ref12
     assert res['definition_refs']['ref2'][0] is ref2
     assert res['recursive_refs'] == set()
-    assert res['deferred_discriminators'] == []
+    assert res['schemas_with_meta_keys'] is None
 
 
 def test_ref_in_serialization_schema():
     ref = core_schema.definition_reference_schema('ref1')
     schema = core_schema.str_schema(
         serialization=core_schema.plain_serializer_function_ser_schema(lambda v: v, return_schema=ref),
     )
-    res = gather_schemas_for_cleaning(schema, definitions={'ref1': core_schema.str_schema()})
+    res = gather_schemas_for_cleaning(schema, definitions={'ref1': core_schema.str_schema()}, find_meta_with_keys=None)
     assert res['definition_refs'] == {'ref1': [ref]} and res['definition_refs']['ref1'][0] is ref
     assert res['recursive_refs'] == set()
-    assert res['deferred_discriminators'] == []
+    assert res['schemas_with_meta_keys'] is None
 
 
 def test_recursive_ref_schema():
     ref1 = core_schema.definition_reference_schema('ref1')
-    res = gather_schemas_for_cleaning(ref1, definitions={'ref1': ref1})
+    res = gather_schemas_for_cleaning(ref1, definitions={'ref1': ref1}, find_meta_with_keys=None)
     assert res['definition_refs'] == {'ref1': [ref1]} and res['definition_refs']['ref1'][0] is ref1
     assert res['recursive_refs'] == {'ref1'}
-    assert res['deferred_discriminators'] == []
+    assert res['schemas_with_meta_keys'] is None
 
 
 def test_deep_recursive_ref_schema():
@@ -78,30 +78,37 @@ def test_deep_recursive_ref_schema():
             'ref2': core_schema.union_schema([ref3, core_schema.float_schema()]),
             'ref3': core_schema.union_schema([ref1, core_schema.str_schema()]),
         },
+        find_meta_with_keys=None,
     )
     assert res['definition_refs'] == {'ref1': [ref1], 'ref2': [ref2], 'ref3': [ref3]}
     assert res['recursive_refs'] == {'ref1', 'ref2', 'ref3'}
     assert res['definition_refs']['ref1'][0] is ref1
     assert res['definition_refs']['ref2'][0] is ref2
     assert res['definition_refs']['ref3'][0] is ref3
-    assert res['deferred_discriminators'] == []
+    assert res['schemas_with_meta_keys'] is None
 
 
-def test_discriminator_meta():
+def test_find_meta():
     class Model:
         pass
 
     ref1 = core_schema.definition_reference_schema('ref1')
 
     field1 = core_schema.model_field(core_schema.str_schema())
-    field1['metadata'] = {'pydantic.internal.union_discriminator': 'foobar1'}
+    field1['metadata'] = {'find_meta1': 'foobar1', 'unknown': 'foobar2'}
 
     field2 = core_schema.model_field(core_schema.int_schema())
-    field2['metadata'] = {'pydantic.internal.union_discriminator': 'foobar2'}
+    field2['metadata'] = {'find_meta1': 'foobar3', 'find_meta2': 'foobar4'}
 
-    schema = core_schema.model_schema(Model, core_schema.model_fields_schema({'a': field1, 'b': ref1}))
-    res = gather_schemas_for_cleaning(schema, definitions={'ref1': field2})
+    schema = core_schema.model_schema(
+        Model, core_schema.model_fields_schema({'a': field1, 'b': ref1, 'c': core_schema.float_schema()})
+    )
+    res = gather_schemas_for_cleaning(
+        schema, definitions={'ref1': field2}, find_meta_with_keys={'find_meta1', 'find_meta2'}
+    )
     assert res['definition_refs'] == {'ref1': [ref1]} and res['definition_refs']['ref1'][0] is ref1
     assert res['recursive_refs'] == set()
-    assert res['deferred_discriminators'] == [(field1, 'foobar1'), (field2, 'foobar2')]
-    assert res['deferred_discriminators'][0][0] is field1 and res['deferred_discriminators'][1][0] is field2
+    assert res['schemas_with_meta_keys'] == {'find_meta1': [field1, field2], 'find_meta2': [field2]}
+    assert res['schemas_with_meta_keys']['find_meta1'][0] is field1
+    assert res['schemas_with_meta_keys']['find_meta1'][1] is field2
+    assert res['schemas_with_meta_keys']['find_meta2'][0] is field2