Enable pickling Bytes (#295)

* Support pickling bytes

* Add test

* Enable hashable

* revert hash

Author: kylebarron

obstore/src/lib.rs

@@ -55,6 +55,9 @@ fn _obstore(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     pyo3_object_store::register_exceptions_module(py, m, "obstore")?;
 
     m.add_class::<pyo3_bytes::PyBytes>()?;
+    // Set the value of `__module__` correctly on PyBytes
+    m.getattr("Bytes")?.setattr("__module__", "obstore")?;
+
     m.add_wrapped(wrap_pyfunction!(buffered::open_reader))?;
     m.add_wrapped(wrap_pyfunction!(buffered::open_reader_async))?;
     m.add_wrapped(wrap_pyfunction!(buffered::open_writer))?;

pyo3-bytes/src/bytes.rs

@@ -8,7 +8,7 @@ use bytes::{Bytes, BytesMut};
 use pyo3::buffer::PyBuffer;
 use pyo3::exceptions::{PyIndexError, PyValueError};
 use pyo3::prelude::*;
-use pyo3::types::PySlice;
+use pyo3::types::{PyDict, PySlice, PyTuple};
 use pyo3::{ffi, IntoPyObjectExt};
 
 /// A wrapper around a [`bytes::Bytes`][].
@@ -161,6 +161,13 @@ impl PyBytes {
         buf
     }
 
+    fn __getnewargs_ex__(&self, py: Python) -> PyResult<PyObject> {
+        let py_bytes = self.to_bytes(py);
+        let args = PyTuple::new(py, vec![py_bytes])?.into_py_any(py)?;
+        let kwargs = PyDict::new(py);
+        PyTuple::new(py, [args, kwargs.into_py_any(py)?])?.into_py_any(py)
+    }
+
     /// The number of bytes in this Bytes
     fn __len__(&self) -> usize {
         self.0.len()

tests/test_bytes.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import pickle
 from typing import TYPE_CHECKING
 
 import pytest
@@ -93,3 +94,8 @@ def _bytes_slices(
             for stop in indices_range
             for step in steps
         )
+
+
+def test_pickle():
+    b = Bytes(b"hello_world")
+    assert b == pickle.loads(pickle.dumps(b))