support pickle/unpickle

aminalaee · Oct 2, 2023 · f11dfba · f11dfba
1 parent 3d5f68c
commit f11dfba
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -75,3 +75,9 @@ UUID('6fa459ea-ee8a-3ca4-894e-db77e160355e')
 
 - The `getnode` function is not available.
 - The `uuid1` and `uuid6` take `node` argument as mandatory.
+
+## How to build
+
+```bash
+RUSTFLAGS="--cfg uuid_unstable" maturin build --release
+```
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,11 +1,7 @@
-use pyo3::{
-    exceptions::{PyTypeError, PyValueError},
-    prelude::*,
-    pyclass::CompareOp,
-    types::PyBytes,
-};
+use pyo3::{exceptions::{PyTypeError, PyValueError}, prelude::*, pyclass::CompareOp, types::PyBytes};
 use std::hash::Hasher;
 use std::{collections::hash_map::DefaultHasher, hash::Hash};
+use pyo3::types::PyTuple;
 use uuid::{Builder, Bytes, Context, Timestamp, Uuid, Variant, Version};
 
 pub const RESERVED_NCS: &str = "reserved for NCS compatibility";
@@ -284,6 +280,10 @@ impl UUID {
             uuid: Uuid::from_u128(int),
         })
     }
+
+    fn __getnewargs__<'a>(&self, py: Python<'a>) -> PyResult<&'a PyTuple> {
+        return Ok(PyTuple::new(py, [self.uuid.to_string()]))
+    }
 }
 
 #[pyfunction]

diff --git a/tests/test_uuid.py b/tests/test_uuid.py
@@ -1,3 +1,4 @@
+import pickle
 from uuid import UUID, getnode
 
 import pytest
@@ -168,3 +169,10 @@ def test_uuid_timestamp() -> None:
 
     with pytest.raises(ValueError):
         uuid_utils.uuid4().timestamp
+
+
+def test_pickle():
+    uuid = uuid_utils.UUID("a8098c1a-f86e-11da-bd1a-00112444be1e")
+    uuid_pickle = pickle.dumps(uuid)
+    uuid_unpickle = pickle.loads(uuid_pickle)
+    assert uuid_unpickle == uuid