[Python] How to add Python pickle to a PyO3 complex Rust enum?

Question


I am using PyO3 as the binding to use Rust from Python. PyO3 gives two categorisations for using #[pyclass] on an enum: simple and complex.

Simple enum seem to be well covered, whereas complex enum have gaps in functionality or accessible solutions on the web. I cannot understand how to use python pickle on a complex enum.

Has anyone been able to implement pickling for a complex enum with PyO3?

Simple Enums can be pickled


When implementing Python pickle for simple enums I generally use the pre-discussed methods of __getnewargs__ or state setting, preferring __getnewargs__ since it doesn't rely on an external choice of serializer.

getnewargs


#[pyclass(eq)]
#[derive(Copy, Clone, PartialEq)]
pub enum MyEnum2 {
Value0 = 0,
Value1 = 1,
}

#[pymethods]
impl MyEnum2 {
// Pickling
#[new]
fn new_py(item: usize) -> MyEnum {
match item {
_ if item == MyEnum2::Value0 as usize => MyEnum2::Value0,
_ if item == MyEnum2::Value1 as usize => MyEnum2::Value1,
_ => MyEnum2::Value0,
}
}
pub fn __getnewargs__<'py>(&self) -> PyResult<(usize,)> {
Ok((*self as usize,))
}
}

state setting


use bincode::config::legacy;
use bincode::serde::{decode_from_slice, encode_to_vec};
// use serde::{Serialize, Deserialize};

#[pymethods]
impl MyEnum2 {
// Pickling
#[new]
fn new_py() -> MyEnum {
MyEnum::Value0
}
// derive Serialize and Deserialize
pub fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
*self = decode_from_slice(state.as_bytes(), legacy()).unwrap().0;
Ok(())
}
pub fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
Ok(PyBytes::new(py, &encode_to_vec(&self, legacy()).unwrap()))
}
}

Complex Enums can't be pickled as above

State setting


#[pyclass(eq)]
#[derive(Debug, Copy, Clone, PartialEq, Serialize, Deserialize)]
pub enum TestEnum {
Value0 {},
Value1 {val: i32},
}

#[pymethods]
impl TestEnum {
// Pickling
#[new]
fn new_py() -> TestEnum {
TestEnum::Value0{}
}
pub fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
*self = decode_from_slice(state.as_bytes(), legacy()).unwrap().0;
Ok(())
}
pub fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
Ok(PyBytes::new(py, &encode_to_vec(&self, legacy()).unwrap()))
}
}


will achieve the error:

error[E0271]: type mismatch resolving `<TestEnum as PyClass>::Frozen == False`

34 | pub fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
| ^ expected `False`, found `True`
104 | pub fn extract_pyclass_ref_mut<'a, 'py: 'a, T: PyClass<Frozen = False>>(
| ^^^^^^^^^^^^^^ required by this bound in `extract_pyclass_ref_mut`

getnewargs


#[pymethods]
impl TestEnum {
// Pickling
#[new]
fn new_py(i: usize, arg: Option<i32>) -> TestEnum {
match i {
0_usize => TestEnum::Value0{},
1_usize => TestEnum::Value1{val: arg.unwrap()},
_ => TestEnum::Value0{}
}
}
pub fn __getnewargs__<'py>(&self) -> PyResult<(usize, Option<i32>)> {
match self {
TestEnum::Value0{} => Ok((0_usize, None)),
TestEnum::Value1{val: n} => Ok((1_usize, Some(*n))),
}
}
}


gets:

>>> pickled = pickle.dumps(TestEnum.Value0())
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E _pickle.PicklingError: Can't pickle <class 'TestEnum_Value0'>: attribute lookup TestEnum_Value0 on module failed.


Anyone able shed any light on this? Perhaps I have made a trivial error also...

Continue reading...