Add SchemaSerializer.__reduce__ method to enable pickle serialization

src/serializers/mod.rs

@@ -3,6 +3,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict};
+use pyo3::PyTypeInfo;
 use pyo3::{PyTraverseError, PyVisit};
 
 use crate::definitions::{Definitions, DefinitionsBuilder};
@@ -33,6 +34,10 @@ pub struct SchemaSerializer {
     definitions: Definitions<CombinedSerializer>,
     expected_json_size: AtomicUsize,
     config: SerializationConfig,
+    // References to the Python schema and config objects are saved to enable
+    // reconstructing the object for cloudpickle support (see `__reduce__`).
+    py_schema: Py<PyDict>,
+    py_config: Option<Py<PyDict>>,
 }
 
 impl SchemaSerializer {
@@ -71,15 +76,19 @@ impl SchemaSerializer {
 #[pymethods]
 impl SchemaSerializer {
     #[new]
-    pub fn py_new(schema: &PyDict, config: Option<&PyDict>) -> PyResult<Self> {
+    pub fn py_new(py: Python, schema: &PyDict, config: Option<&PyDict>) -> PyResult<Self> {
         let mut definitions_builder = DefinitionsBuilder::new();
-
         let serializer = CombinedSerializer::build(schema.downcast()?, config, &mut definitions_builder)?;
         Ok(Self {
             serializer,
             definitions: definitions_builder.finish()?,
             expected_json_size: AtomicUsize::new(1024),
             config: SerializationConfig::from_config(config)?,
+            py_schema: schema.into_py(py),
+            py_config: match config {
+                Some(d) if !d.is_empty() => Some(d.into_py(py)),
+                _ => None,
+            },
         })
     }
 
@@ -174,6 +183,14 @@ impl SchemaSerializer {
         Ok(py_bytes.into())
     }
 
+    pub fn __reduce__(&self, py: Python) -> PyResult<(PyObject, (PyObject, PyObject))> {
+        // Enables support for `cloudpickle` serialization.
+        Ok((
+            SchemaSerializer::type_object(py).to_object(py),
+            (self.py_schema.to_object(py), self.py_config.to_object(py)),
+        ))
+    }
+
     pub fn __repr__(&self) -> String {
         format!(
             "SchemaSerializer(serializer={:#?}, definitions={:#?})",
@@ -182,6 +199,10 @@ impl SchemaSerializer {
     }
 
     fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.py_schema)?;
+        if let Some(ref py_config) = self.py_config {
+            visit.call(py_config)?;
+        }
         self.serializer.py_gc_traverse(&visit)?;
         self.definitions.py_gc_traverse(&visit)?;
         Ok(())

src/validators/url.rs

@@ -498,7 +498,7 @@ fn check_sub_defaults(
         if let Some(default_port) = default_port {
             lib_url
                 .set_port(Some(default_port))
-                .map_err(|_| map_parse_err(ParseError::EmptyHost))?;
+                .map_err(|()| map_parse_err(ParseError::EmptyHost))?;
         }
     }
     if let Some(ref default_path) = default_path {

tests/test.rs

@@ -46,7 +46,7 @@ mod tests {
                 ]
             }"#;
             let schema: &PyDict = py.eval(code, None, None).unwrap().extract().unwrap();
-            SchemaSerializer::py_new(schema, None).unwrap();
+            SchemaSerializer::py_new(py, schema, None).unwrap();
         });
     }
 
@@ -77,7 +77,7 @@ a = A()
             py.run(code, None, Some(locals)).unwrap();
             let a: &PyAny = locals.get_item("a").unwrap().extract().unwrap();
             let schema: &PyDict = locals.get_item("schema").unwrap().extract().unwrap();
-            let serialized: Vec<u8> = SchemaSerializer::py_new(schema, None)
+            let serialized: Vec<u8> = SchemaSerializer::py_new(py, schema, None)
                 .unwrap()
                 .to_json(py, a, None, None, None, true, false, false, false, false, true, None)
                 .unwrap()

tests/test_garbage_collection.py

@@ -27,7 +27,9 @@ class BaseModel:
         __schema__: SchemaSerializer
 
         def __init_subclass__(cls) -> None:
-            cls.__schema__ = SchemaSerializer(core_schema.model_schema(cls, GC_TEST_SCHEMA_INNER))
+            cls.__schema__ = SchemaSerializer(
+                core_schema.model_schema(cls, GC_TEST_SCHEMA_INNER), config={'ser_json_timedelta': 'float'}
+            )
 
     cache: 'WeakValueDictionary[int, Any]' = WeakValueDictionary()
 

tests/test_pickling.py

@@ -0,0 +1,50 @@
+import json
+import pickle
+from datetime import timedelta
+
+import pytest
+
+from pydantic_core import core_schema
+from pydantic_core._pydantic_core import SchemaSerializer
+
+
+def repr_function(value, _info):
+    return repr(value)
+
+
+def test_basic_schema_serializer():
+    s = SchemaSerializer(core_schema.dict_schema())
+    s = pickle.loads(pickle.dumps(s))
+    assert s.to_python({'a': 1, b'b': 2, 33: 3}) == {'a': 1, b'b': 2, 33: 3}
+    assert s.to_python({'a': 1, b'b': 2, 33: 3, True: 4}, mode='json') == {'a': 1, 'b': 2, '33': 3, 'true': 4}
+    assert s.to_json({'a': 1, b'b': 2, 33: 3, True: 4}) == b'{"a":1,"b":2,"33":3,"true":4}'
+
+    assert s.to_python({(1, 2): 3}) == {(1, 2): 3}
+    assert s.to_python({(1, 2): 3}, mode='json') == {'1,2': 3}
+    assert s.to_json({(1, 2): 3}) == b'{"1,2":3}'
+
+
+@pytest.mark.parametrize(
+    'value,expected_python,expected_json',
+    [(None, 'None', b'"None"'), (1, '1', b'"1"'), ([1, 2, 3], '[1, 2, 3]', b'"[1, 2, 3]"')],
+)
+def test_schema_serializer_capturing_function(value, expected_python, expected_json):
+    # Test a SchemaSerializer that captures a function.
+    s = SchemaSerializer(
+        core_schema.any_schema(
+            serialization=core_schema.plain_serializer_function_ser_schema(repr_function, info_arg=True)
+        )
+    )
+    s = pickle.loads(pickle.dumps(s))
+    assert s.to_python(value) == expected_python
+    assert s.to_json(value) == expected_json
+    assert s.to_python(value, mode='json') == json.loads(expected_json)
+
+
+def test_schema_serializer_containing_config():
+    s = SchemaSerializer(core_schema.timedelta_schema(), config={'ser_json_timedelta': 'float'})
+    s = pickle.loads(pickle.dumps(s))
+
+    assert s.to_python(timedelta(seconds=4, microseconds=500_000)) == timedelta(seconds=4, microseconds=500_000)
+    assert s.to_python(timedelta(seconds=4, microseconds=500_000), mode='json') == 4.5
+    assert s.to_json(timedelta(seconds=4, microseconds=500_000)) == b'4.5'