Improve res type

Author: MarkusSintonen

python/pydantic_core/__init__.py

@@ -2,7 +2,6 @@
 
 import sys as _sys
 from typing import Any as _Any
-from typing import NamedTuple
 
 from ._pydantic_core import (
     ArgsKwargs,
@@ -142,7 +141,9 @@ class MultiHostHost(_TypedDict):
     """The port part of this host, or `None`."""
 
 
-class GatherResult(NamedTuple):
+class GatherResult(_TypedDict):
+    """Internal result of gathering schemas for cleaning."""
+
     definition_refs: dict[str, list[DefinitionReferenceSchema]]
     recursive_refs: set[str]
     deferred_discriminators: list[tuple[CoreSchema, _Any]]

src/lib.rs

@@ -36,11 +36,12 @@ 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 serializers::{
     to_json, to_jsonable_python, PydanticSerializationError, PydanticSerializationUnexpectedValue, SchemaSerializer,
     WarningsArg,
 };
-pub use validators::{gather_schemas_for_cleaning, validate_core_schema, PySome, SchemaValidator};
+pub use validators::{validate_core_schema, PySome, SchemaValidator};
 
 use crate::input::Input;
 

src/schema_traverse.rs

@@ -202,7 +202,25 @@ impl<'a, 'py> GatherCtx<'a, 'py> {
     }
 }
 
-pub fn gather_schema(schema: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()> {
+fn gather_schema(schema: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()> {
     traverse_schema(schema, ctx)?;
     gather_node(schema, ctx)
 }
+
+#[pyfunction(signature = (schema, definitions))]
+pub fn gather_schemas_for_cleaning<'py>(
+    schema: &Bound<'py, PyAny>,
+    definitions: &Bound<'py, PyAny>,
+) -> PyResult<Bound<'py, PyDict>> {
+    let py = schema.py();
+    let schema_dict = schema.downcast_exact::<PyDict>()?;
+
+    let mut ctx = GatherCtx::new(definitions.downcast_exact()?)?;
+    gather_schema(schema_dict, &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)?;
+    Ok(res)
+}

src/validators/mod.rs

@@ -64,7 +64,6 @@ mod validation_state;
 mod with_default;
 
 pub use self::validation_state::{Exactness, ValidationState};
-use crate::schema_traverse::{gather_schema, GatherCtx};
 pub use with_default::DefaultType;
 
 #[pyclass(module = "pydantic_core._pydantic_core", name = "Some")]
@@ -444,25 +443,6 @@ impl<'py> SelfValidator<'py> {
     }
 }
 
-#[pyfunction(signature = (schema, definitions))]
-pub fn gather_schemas_for_cleaning<'py>(
-    schema: &Bound<'py, PyAny>,
-    definitions: &Bound<'py, PyAny>,
-) -> PyResult<Bound<'py, PyTuple>> {
-    let py = schema.py();
-    let schema_dict = schema.downcast_exact::<PyDict>()?;
-
-    let mut ctx = GatherCtx::new(definitions.downcast_exact()?)?;
-    gather_schema(schema_dict, &mut ctx)?;
-
-    let res = vec![
-        ctx.def_refs.as_any(),
-        ctx.recursive_def_refs.as_any(),
-        ctx.discriminators.as_any(),
-    ];
-    return Ok(PyTuple::new_bound(py, res));
-}
-
 #[pyfunction(signature = (schema, *, strict = None))]
 pub fn validate_core_schema<'py>(schema: &Bound<'py, PyAny>, strict: Option<bool>) -> PyResult<Bound<'py, PyAny>> {
     let self_validator = SelfValidator::new(schema.py())?;