fnv pickling

jessekrubin · Mar 10, 2025 · 52b6579 · 52b6579
1 parent 7886cd6
commit 52b6579
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 41 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/ryo3-fnv/src/lib.rs b/crates/ryo3-fnv/src/lib.rs
@@ -2,8 +2,9 @@
 use std::hash::Hasher;
 
 use ::fnv as fnv_rs;
-use pyo3::types::PyModule;
-use pyo3::{intern, prelude::*};
+use pyo3::types::{PyModule, PyTuple};
+
+use pyo3::{intern, prelude::*, IntoPyObjectExt};
 use pyo3::{wrap_pyfunction, PyResult};
 
 #[pyclass(name = "FnvHasher", module = "ryo3")]
@@ -14,20 +15,38 @@ pub struct PyFnvHasher {
 #[pymethods]
 impl PyFnvHasher {
     #[new]
-    #[pyo3(signature = (s = None))]
-    fn py_new(s: Option<&[u8]>) -> Self {
-        match s {
-            Some(s) => {
+    #[pyo3(signature = (s = None, *, key = None))]
+    fn py_new(s: Option<&[u8]>, key: Option<u64>) -> Self {
+        match (key, s) {
+            (Some(k), Some(s)) => {
+                let mut hasher = fnv_rs::FnvHasher::with_key(k);
+                hasher.write(s);
+                Self { hasher }
+            }
+            (Some(k), None) => Self {
+                hasher: fnv_rs::FnvHasher::with_key(k),
+            },
+            (None, Some(s)) => {
                 let mut hasher = fnv_rs::FnvHasher::default();
                 hasher.write(s);
                 Self { hasher }
             }
-            None => Self {
+            (None, None) => Self {
                 hasher: fnv_rs::FnvHasher::default(),
             },
         }
     }
 
+    fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        PyTuple::new(
+            py,
+            [
+                py.None().into_bound_py_any(py)?,
+                self.hasher.finish().into_bound_py_any(py)?,
+            ],
+        )
+    }
+
     fn __str__(&self) -> String {
         format!("fnv1a<{:x}>", self.hasher.finish())
     }
@@ -63,8 +82,19 @@ impl PyFnvHasher {
 }
 
 #[pyfunction]
-pub fn fnv1a(s: &[u8]) -> PyResult<PyFnvHasher> {
-    Ok(PyFnvHasher::py_new(Some(s)))
+#[pyo3(signature = (s, key = None))]
+pub fn fnv1a(s: &[u8], key: Option<u64>) -> PyResult<PyFnvHasher> {
+    Ok(PyFnvHasher {
+        hasher: if let Some(k) = key {
+            let mut hasher = fnv_rs::FnvHasher::with_key(k);
+            hasher.write(s);
+            hasher
+        } else {
+            let mut hasher = fnv_rs::FnvHasher::default();
+            hasher.write(s);
+            hasher
+        },
+    })
 }
 
 pub fn pymod_add(m: &Bound<'_, PyModule>) -> PyResult<()> {