Add exception for unknown schema_ref

Author: MarkusSintonen

python/pydantic_core/__init__.py

@@ -5,6 +5,7 @@
 
 from ._pydantic_core import (
     ArgsKwargs,
+    GatherInvalidDefinitionError,
     MultiHostUrl,
     PydanticCustomError,
     PydanticKnownError,
@@ -63,6 +64,7 @@
     'PydanticUseDefault',
     'PydanticSerializationError',
     'PydanticSerializationUnexpectedValue',
+    'GatherInvalidDefinitionError',
     'TzInfo',
     'to_json',
     'from_json',

python/pydantic_core/_pydantic_core.pyi

@@ -28,6 +28,7 @@ __all__ = [
     'PydanticSerializationUnexpectedValue',
     'PydanticUndefined',
     'PydanticUndefinedType',
+    'GatherInvalidDefinitionError',
     'Some',
     'to_json',
     'from_json',
@@ -1165,6 +1166,9 @@ def validate_core_schema(schema: CoreSchema, *, strict: bool | None = None) -> C
     We may also remove this function altogether, do not rely on it being present if you are
     using pydantic-core directly.
     """
+@final
+class GatherInvalidDefinitionError(Exception):
+    """Internal error when encountering invalid definition refs"""
 
 def gather_schemas_for_cleaning(
     schema: CoreSchema,

src/lib.rs

@@ -36,7 +36,7 @@ pub use build_tools::SchemaError;
 pub use errors::{
     list_all_errors, PydanticCustomError, PydanticKnownError, PydanticOmit, PydanticUseDefault, ValidationError,
 };
-pub use schema_traverse::gather_schemas_for_cleaning;
+pub use schema_traverse::{gather_schemas_for_cleaning, GatherInvalidDefinitionError};
 pub use serializers::{
     to_json, to_jsonable_python, PydanticSerializationError, PydanticSerializationUnexpectedValue, SchemaSerializer,
     WarningsArg,
@@ -131,6 +131,10 @@ fn _pydantic_core(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<ArgsKwargs>()?;
     m.add_class::<SchemaSerializer>()?;
     m.add_class::<TzInfo>()?;
+    m.add(
+        "GatherInvalidDefinitionError",
+        py.get_type_bound::<GatherInvalidDefinitionError>(),
+    )?;
     m.add_function(wrap_pyfunction!(to_json, m)?)?;
     m.add_function(wrap_pyfunction!(from_json, m)?)?;
     m.add_function(wrap_pyfunction!(to_jsonable_python, m)?)?;

src/schema_traverse.rs

@@ -1,10 +1,12 @@
 use crate::tools::py_err;
-use pyo3::exceptions::PyKeyError;
+use pyo3::exceptions::{PyException, PyKeyError};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PySet, PyString, PyTuple};
-use pyo3::{intern, Bound, PyResult};
+use pyo3::{create_exception, intern, Bound, PyResult};
 use std::collections::HashSet;
 
+create_exception!(pydantic_core._pydantic_core, GatherInvalidDefinitionError, PyException);
+
 macro_rules! get {
     ($dict: expr, $key: expr) => {
         $dict.get_item(intern!($dict.py(), $key))?
@@ -42,44 +44,47 @@ macro_rules! defaultdict_list_append {
 }
 
 fn gather_definition_ref(schema_ref_dict: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()> {
-    if let Some(schema_ref) = get!(schema_ref_dict, "schema_ref") {
-        let schema_ref_pystr = schema_ref.downcast_exact::<PyString>()?;
-        let schema_ref_str = schema_ref_pystr.to_str()?;
+    let Some(schema_ref) = get!(schema_ref_dict, "schema_ref") else {
+        return py_err!(PyKeyError; "Invalid definition-ref, missing schema_ref");
+    };
+    let schema_ref_pystr = schema_ref.downcast_exact::<PyString>()?;
+    let schema_ref_str = schema_ref_pystr.to_str()?;
 
-        if !ctx.recursively_seen_refs.contains(schema_ref_str) {
-            defaultdict_list_append!(&ctx.def_refs, schema_ref_pystr, schema_ref_dict);
+    if !ctx.recursively_seen_refs.contains(schema_ref_str) {
+        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.get_item(schema_ref_pystr)? {
-                ctx.recursively_seen_refs.insert(schema_ref_str.to_string());
+        let Some(definition) = ctx.definitions.get_item(schema_ref_pystr)? else {
+            return py_err!(GatherInvalidDefinitionError; "Unknown schema_ref: {}", schema_ref_str);
+        };
 
-                gather_schema(definition.downcast_exact::<PyDict>()?, ctx)?;
-                traverse_key_fn!("serialization", gather_schema, schema_ref_dict, ctx);
-                gather_meta(schema_ref_dict, ctx)?;
+        ctx.recursively_seen_refs.insert(schema_ref_str.to_string());
 
-                ctx.recursively_seen_refs.remove(schema_ref_str);
-            }
-        } else {
-            ctx.recursive_def_refs.add(schema_ref_pystr)?;
-            for seen_ref in &ctx.recursively_seen_refs {
-                let seen_ref_pystr = PyString::new_bound(schema_ref.py(), seen_ref);
-                ctx.recursive_def_refs.add(seen_ref_pystr)?;
-            }
-        }
-        Ok(())
+        gather_schema(definition.downcast_exact()?, ctx)?;
+        traverse_key_fn!("serialization", gather_schema, schema_ref_dict, ctx);
+        gather_meta(schema_ref_dict, ctx)?;
+
+        ctx.recursively_seen_refs.remove(schema_ref_str);
     } else {
-        py_err!(PyKeyError; "Invalid definition-ref, missing schema_ref")
+        ctx.recursive_def_refs.add(schema_ref_pystr)?;
+        for seen_ref in &ctx.recursively_seen_refs {
+            let seen_ref_pystr = PyString::new_bound(schema_ref.py(), seen_ref);
+            ctx.recursive_def_refs.add(seen_ref_pystr)?;
+        }
     }
+    Ok(())
 }
 
 fn gather_meta(schema: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()> {
-    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);
-                }
-            }
+    let Some((res, find_keys)) = &ctx.meta_with_keys else {
+        return Ok(());
+    };
+    let Some(meta) = get!(schema, "metadata") else {
+        return Ok(());
+    };
+    let meta_dict = meta.downcast_exact::<PyDict>()?;
+    for k in find_keys.iter() {
+        if meta_dict.contains(&k)? {
+            defaultdict_list_append!(res, &k, schema);
         }
     }
     Ok(())
@@ -120,11 +125,10 @@ fn gather_arguments(arguments: &Bound<'_, PyList>, ctx: &mut GatherCtx) -> PyRes
 // Has 100% coverage in Pydantic side. This is exclusively used there
 #[cfg_attr(has_coverage_attribute, coverage(off))]
 fn gather_schema(schema: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()> {
-    let type_ = get!(schema, "type");
-    if type_.is_none() {
+    let Some(type_) = get!(schema, "type") else {
         return py_err!(PyKeyError; "Schema type missing");
-    }
-    match type_.unwrap().downcast_exact::<PyString>()?.to_str()? {
+    };
+    match type_.downcast_exact::<PyString>()?.to_str()? {
         "definition-ref" => gather_definition_ref(schema, ctx),
         "definitions" => traverse!("schema" => gather_schema, "definitions" => gather_list; schema, ctx),
         "list" | "set" | "frozenset" | "generator" => traverse!("items_schema" => gather_schema; schema, ctx),
@@ -167,14 +171,12 @@ pub fn gather_schemas_for_cleaning<'py>(
     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: definitions.downcast_exact()?,
-        meta_with_keys,
+        meta_with_keys: match find_meta_with_keys.is_none() {
+            true => None,
+            false => Some((PyDict::new_bound(py), find_meta_with_keys.downcast_exact::<PySet>()?)),
+        },
         def_refs: PyDict::new_bound(py),
         recursive_def_refs: PySet::empty_bound(py)?,
         recursively_seen_refs: HashSet::new(),

tests/test_gather_schemas_for_cleaning.py

@@ -1,4 +1,6 @@
-from pydantic_core import core_schema, gather_schemas_for_cleaning
+import pytest
+
+from pydantic_core import GatherInvalidDefinitionError, core_schema, gather_schemas_for_cleaning
 
 
 def test_no_refs():
@@ -112,3 +114,10 @@ class Model:
     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
+
+
+def test_unknown_ref():
+    ref1 = core_schema.definition_reference_schema('ref1')
+    schema = core_schema.tuple_schema([core_schema.int_schema(), ref1])
+    with pytest.raises(GatherInvalidDefinitionError, match='Unknown schema_ref: ref1'):
+        gather_schemas_for_cleaning(schema, definitions={}, find_meta_with_keys=None)