make Url and MultiHostUrl pickleable (pydantic#676)

alexdrydew · Jun 21, 2023 · 0ff7563 · 0ff7563
1 parent 4420be3
commit 0ff7563
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/src/url.rs b/src/url.rs
@@ -145,6 +145,10 @@ impl PyUrl {
     pub fn __deepcopy__(&self, py: Python, _memo: &PyDict) -> Py<PyAny> {
         self.clone().into_py(py)
     }
+
+    fn __getnewargs__(&self) -> (&str,) {
+        (self.__str__(),)
+    }
 }
 
 #[pyclass(name = "MultiHostUrl", module = "pydantic_core._pydantic_core", subclass)]
@@ -305,6 +309,10 @@ impl PyMultiHostUrl {
     pub fn __deepcopy__(&self, py: Python, _memo: &PyDict) -> Py<PyAny> {
         self.clone().into_py(py)
     }
+
+    fn __getnewargs__(&self) -> (String,) {
+        (self.__str__(),)
+    }
 }
 
 fn host_to_dict<'a>(py: Python<'a>, lib_url: &Url) -> PyResult<&'a PyDict> {

diff --git a/tests/serializers/test_url.py b/tests/serializers/test_url.py
@@ -1,3 +1,5 @@
+import pickle
+
 import pytest
 
 from pydantic_core import MultiHostUrl, SchemaSerializer, SchemaValidator, Url, core_schema
@@ -96,3 +98,10 @@ def some_method(self):
 
     m = MyUrl('http://ex.com/path')
     assert m.some_method() == '/path-success'
+
+
+@pytest.mark.parametrize('value', (Url('https://example.com'), MultiHostUrl('https://example.com,example.org/path')))
+def test_url_pickle(value):
+    pickled = pickle.dumps(value)
+    unpickled = pickle.loads(pickled)
+    assert value == unpickled