fix python GC traversal for validators and serializers

src/lib.rs

@@ -10,6 +10,10 @@ use pyo3::{prelude::*, sync::GILOnceCell};
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
+// parse this first to get access to the contained macro
+#[macro_use]
+mod py_gc;
+
 mod argument_markers;
 mod build_tools;
 mod definitions;

src/py_gc.rs

@@ -0,0 +1,70 @@
+use ahash::AHashMap;
+use enum_dispatch::enum_dispatch;
+use pyo3::{AsPyPointer, Py, PyTraverseError, PyVisit};
+
+/// Trait implemented by types which can be traversed by the Python GC.
+#[enum_dispatch]
+pub trait PyGcTraverse {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
+}
+
+impl<T> PyGcTraverse for Py<T>
+where
+    Py<T>: AsPyPointer,
+{
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(self)
+    }
+}
+
+impl<T: PyGcTraverse> PyGcTraverse for Vec<T> {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        for item in self {
+            item.py_gc_traverse(visit)?;
+        }
+        Ok(())
+    }
+}
+
+impl<T: PyGcTraverse> PyGcTraverse for AHashMap<String, T> {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        for item in self.values() {
+            item.py_gc_traverse(visit)?;
+        }
+        Ok(())
+    }
+}
+
+impl<T: PyGcTraverse> PyGcTraverse for Box<T> {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        T::py_gc_traverse(self, visit)
+    }
+}
+
+impl<T: PyGcTraverse> PyGcTraverse for Option<T> {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        match self {
+            Some(item) => T::py_gc_traverse(item, visit),
+            None => Ok(()),
+        }
+    }
+}
+
+/// A crude alternative to a "derive" macro to help with building PyGcTraverse implementations
+macro_rules! impl_py_gc_traverse {
+    ($name:ty { }) => {
+        impl crate::py_gc::PyGcTraverse for $name {
+            fn py_gc_traverse(&self, _visit: &pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {
+                Ok(())
+            }
+        }
+    };
+    ($name:ty { $($fields:ident),* }) => {
+        impl crate::py_gc::PyGcTraverse for $name {
+            fn py_gc_traverse(&self, visit: &pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {
+                $(self.$fields.py_gc_traverse(visit)?;)*
+                Ok(())
+            }
+        }
+    };
+}

src/serializers/computed_fields.rs

@@ -1,11 +1,12 @@
-use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyString};
+use pyo3::{intern, PyTraverseError, PyVisit};
 use serde::ser::SerializeMap;
 use serde::Serialize;
 
 use crate::build_tools::py_schema_error_type;
 use crate::definitions::DefinitionsBuilder;
+use crate::py_gc::PyGcTraverse;
 use crate::serializers::filter::SchemaFilter;
 use crate::serializers::shared::{BuildSerializer, CombinedSerializer, PydanticSerializer, TypeSerializer};
 use crate::tools::SchemaDict;
@@ -156,6 +157,16 @@ pub(crate) struct ComputedFieldSerializer<'py> {
     extra: &'py Extra<'py>,
 }
 
+impl_py_gc_traverse!(ComputedField { serializer });
+
+impl PyGcTraverse for ComputedFields {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        self.0.py_gc_traverse(visit)
+    }
+}
+
+impl_py_gc_traverse!(ComputedFieldSerializer<'_> { computed_field });
+
 impl<'py> Serialize for ComputedFieldSerializer<'py> {
     fn serialize<S: serde::ser::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         let py = self.model.py();

src/serializers/fields.rs

@@ -28,6 +28,8 @@ pub(super) struct SerField {
     pub required: bool,
 }
 
+impl_py_gc_traverse!(SerField { serializer });
+
 impl SerField {
     pub fn new(
         py: Python,
@@ -142,6 +144,11 @@ macro_rules! option_length {
     };
 }
 
+impl_py_gc_traverse!(GeneralFieldsSerializer {
+    fields,
+    computed_fields
+});
+
 impl TypeSerializer for GeneralFieldsSerializer {
     fn to_python(
         &self,

src/serializers/mod.rs

@@ -6,6 +6,7 @@ use pyo3::types::{PyBytes, PyDict};
 use pyo3::{PyTraverseError, PyVisit};
 
 use crate::definitions::DefinitionsBuilder;
+use crate::py_gc::PyGcTraverse;
 use crate::validators::SelfValidator;
 
 use config::SerializationConfig;
@@ -191,13 +192,6 @@ impl SchemaSerializer {
         }
         Ok(())
     }
-
-    fn __clear__(&mut self) {
-        self.serializer.py_gc_clear();
-        for slot in &mut self.definitions {
-            slot.py_gc_clear();
-        }
-    }
 }
 
 #[allow(clippy::too_many_arguments)]

src/serializers/shared.rs

@@ -14,6 +14,7 @@ use serde_json::ser::PrettyFormatter;
 use crate::build_tools::py_schema_err;
 use crate::build_tools::py_schema_error_type;
 use crate::definitions::DefinitionsBuilder;
+use crate::py_gc::PyGcTraverse;
 use crate::tools::{py_err, SchemaDict};
 
 use super::errors::se_err_py_err;
@@ -215,12 +216,50 @@ impl BuildSerializer for CombinedSerializer {
     }
 }
 
+// Implemented by hand because `enum_dispatch` fails with a proc macro compile error =/
+impl PyGcTraverse for CombinedSerializer {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        match self {
+            CombinedSerializer::Function(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::FunctionWrap(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Fields(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::None(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Nullable(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Int(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Bool(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Float(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Str(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Bytes(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Datetime(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::TimeDelta(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Date(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Time(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::List(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Set(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::FrozenSet(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Generator(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Dict(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Model(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Dataclass(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Url(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::MultiHostUrl(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Any(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Format(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::ToString(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::WithDefault(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Json(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::JsonOrPython(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Union(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Literal(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::Recursive(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::TuplePositional(inner) => inner.py_gc_traverse(visit),
+            CombinedSerializer::TupleVariable(inner) => inner.py_gc_traverse(visit),
+        }
+    }
+}
+
 #[enum_dispatch(CombinedSerializer)]
 pub(crate) trait TypeSerializer: Send + Sync + Clone + Debug {
-    fn py_gc_traverse(&self, _visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        Ok(())
-    }
-    fn py_gc_clear(&mut self) {}
     fn to_python(
         &self,
         value: &PyAny,

src/serializers/type_serializers/any.rs

@@ -26,6 +26,8 @@ impl BuildSerializer for AnySerializer {
     }
 }
 
+impl_py_gc_traverse!(AnySerializer {});
+
 impl TypeSerializer for AnySerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/bytes.rs

@@ -25,6 +25,8 @@ impl BuildSerializer for BytesSerializer {
     }
 }
 
+impl_py_gc_traverse!(BytesSerializer {});
+
 impl TypeSerializer for BytesSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/dataclass.rs

@@ -1,6 +1,6 @@
+use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyString, PyType};
-use pyo3::{intern, PyTraverseError, PyVisit};
 use std::borrow::Cow;
 
 use ahash::AHashMap;
@@ -121,13 +121,9 @@ impl DataclassSerializer {
     }
 }
 
-impl TypeSerializer for DataclassSerializer {
-    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
-        visit.call(&self.class)?;
-        self.serializer.py_gc_traverse(visit)?;
-        Ok(())
-    }
+impl_py_gc_traverse!(DataclassSerializer { class, serializer });
 
+impl TypeSerializer for DataclassSerializer {
     fn to_python(
         &self,
         value: &PyAny,

src/serializers/type_serializers/datetime_etc.rs

@@ -40,6 +40,8 @@ macro_rules! build_serializer {
             }
         }
 
+        impl_py_gc_traverse!($struct_name {});
+
         impl TypeSerializer for $struct_name {
             fn to_python(
                 &self,

src/serializers/type_serializers/definitions.rs

@@ -5,6 +5,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList};
 
 use crate::definitions::DefinitionsBuilder;
+
 use crate::tools::SchemaDict;
 
 use super::{py_err_se_err, BuildSerializer, CombinedSerializer, Extra, TypeSerializer};
@@ -59,6 +60,8 @@ impl BuildSerializer for DefinitionRefSerializer {
     }
 }
 
+impl_py_gc_traverse!(DefinitionRefSerializer {});
+
 impl TypeSerializer for DefinitionRefSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/dict.rs

@@ -65,6 +65,11 @@ impl BuildSerializer for DictSerializer {
     }
 }
 
+impl_py_gc_traverse!(DictSerializer {
+    key_serializer,
+    value_serializer
+});
+
 impl TypeSerializer for DictSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/format.rs

@@ -104,6 +104,8 @@ impl FormatSerializer {
     }
 }
 
+impl_py_gc_traverse!(FormatSerializer { format_func });
+
 impl TypeSerializer for FormatSerializer {
     fn to_python(
         &self,
@@ -175,6 +177,8 @@ impl BuildSerializer for ToStringSerializer {
     }
 }
 
+impl_py_gc_traverse!(ToStringSerializer {});
+
 impl TypeSerializer for ToStringSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/function.rs

@@ -269,6 +269,12 @@ macro_rules! function_type_serializer {
     };
 }
 
+impl_py_gc_traverse!(FunctionPlainSerializer {
+    func,
+    return_serializer,
+    fallback_serializer
+});
+
 function_type_serializer!(FunctionPlainSerializer);
 
 fn copy_outer_schema(schema: &PyDict) -> PyResult<&PyDict> {
@@ -399,6 +405,12 @@ impl FunctionWrapSerializer {
     }
 }
 
+impl_py_gc_traverse!(FunctionWrapSerializer {
+    serializer,
+    func,
+    return_serializer
+});
+
 function_type_serializer!(FunctionWrapSerializer);
 
 #[pyclass(module = "pydantic_core._pydantic_core")]

src/serializers/type_serializers/generator.rs

@@ -42,6 +42,8 @@ impl BuildSerializer for GeneratorSerializer {
     }
 }
 
+impl_py_gc_traverse!(GeneratorSerializer { item_serializer });
+
 impl TypeSerializer for GeneratorSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/json.rs

@@ -42,6 +42,8 @@ impl BuildSerializer for JsonSerializer {
     }
 }
 
+impl_py_gc_traverse!(JsonSerializer { serializer });
+
 impl TypeSerializer for JsonSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/json_or_python.rs

@@ -45,6 +45,8 @@ impl BuildSerializer for JsonOrPythonSerializer {
     }
 }
 
+impl_py_gc_traverse!(JsonOrPythonSerializer { json, python });
+
 impl TypeSerializer for JsonOrPythonSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/list.rs

@@ -45,6 +45,8 @@ impl BuildSerializer for ListSerializer {
     }
 }
 
+impl_py_gc_traverse!(ListSerializer { item_serializer });
+
 impl TypeSerializer for ListSerializer {
     fn to_python(
         &self,

src/serializers/type_serializers/literal.rs

@@ -106,6 +106,8 @@ impl LiteralSerializer {
     }
 }
 
+impl_py_gc_traverse!(LiteralSerializer { expected_py });
+
 impl TypeSerializer for LiteralSerializer {
     fn to_python(
         &self,