Skip to content

Instantly share code, notes, and snippets.

@trungnt13
Forked from ethanhs/example.rs
Created April 21, 2024 11:53
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save trungnt13/539243aa8155a7ad58a58baf633640bd to your computer and use it in GitHub Desktop.
Save trungnt13/539243aa8155a7ad58a58baf633640bd to your computer and use it in GitHub Desktop.
Example using pickling in pyo3
#![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