PyO3 pickle for complex Enums

1

To implement Python pickle support for simple Rust Enums using PyO3 I can adopt one of two possibilities:

#[pyclass(eq)]
#[derive(Copy, Clone, PartialEq)]
pub enum MyEnum {
    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,))
    }
}

or

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()))
    }
}

However if my Enum is complex then using the state implementation method:

/// A roll day.
#[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()))
    }
}

I receive the compilation 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`

So I'm assuming that this only works for frozen classes and this complex Enum is not frozen so it won't work.

Simultaneously, if I adopt the new args approach:

#[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))),
        }
    }
}

then I will attain the error in Python:

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

Can anyone suggest a simple solution for this case?