-
-
Save trungnt13/539243aa8155a7ad58a58baf633640bd to your computer and use it in GitHub Desktop.
Example using pickling in pyo3
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#![feature(arbitrary_self_types)] | |
use pyo3::prelude::*; | |
use pyo3::pyclass::PyClassShell; | |
use pyo3::types::{PyBytes, PyTuple}; | |
use pyo3::ToPyObject; | |
use bincode::{deserialize, serialize}; | |
use serde::{Deserialize, Serialize}; | |
#[derive(Serialize, Deserialize)] | |
struct ARustThing { | |
f: usize, | |
} | |
#[pyclass(module = "pickleme")] | |
struct PickleMe { | |
attr: Option<ARustThing>, | |
#[pyo3(get, set)] | |
foo: Option<usize>, | |
} | |
#[pymethods] | |
impl PickleMe { | |
#[new] | |
#[args(args = "*")] | |
fn new(args: &PyTuple) -> Self { | |
match args.len() { | |
0 => PickleMe { | |
attr: None, | |
foo: None, | |
}, | |
1 => { | |
if let Ok(f) = args.get_item(0).extract::<usize>() { | |
PickleMe { | |
attr: Some(ARustThing { f }), | |
foo: None, | |
} | |
} else { | |
PickleMe { | |
attr: None, | |
foo: None, | |
} | |
} | |
} | |
_ => unreachable!(), | |
} | |
} | |
pub fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> { | |
match state.extract::<&PyBytes>(py) { | |
Ok(s) => { | |
self.foo = deserialize(s.as_bytes()).unwrap(); | |
Ok(()) | |
} | |
Err(e) => Err(e), | |
} | |
} | |
pub fn __getstate__(&self, py: Python) -> PyResult<PyObject> { | |
Ok(PyBytes::new(py, &serialize(&self.foo).unwrap()).to_object(py)) | |
} | |
// __reduce__ is nice as it is an (essentially) all in one thing, you pass a tuple of (type, (arguments to type.__new__)) | |
/* | |
pub fn __reduce__(self: &PyClassShell<Self>) -> PyResult<(PyObject, PyObject)> { | |
let gil = Python::acquire_gil(); | |
let py = gil.python(); | |
let cls = self.to_object(py).getattr(py, "__class__")?; | |
Ok(( | |
cls, | |
(PyBytes::new(py, &serialize(&self.foo).unwrap()).to_object(py),).to_object(py), | |
)) | |
} | |
*/ | |
} | |
#[pymodule] | |
fn pickleme(_py: Python, m: &PyModule) -> PyResult<()> { | |
m.add_class::<PickleMe>()?; | |
Ok(()) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment