How to Pickle a Complex Enum?

[attack68]

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?

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...

[attack68]

OK, after weeks of fumbling around with unfortunately little pre-existing information to go on I have managed to get something working.

Question

Objective

I have the following complex enum defined in Rust and ported to Python with PyO3 and I want it to be pickleable.

#[pyclass(module = "mymodule.rs", eq)]
#[derive(Clone, Debug, PartialEq)]
pub enum MyEnum {
    Single(i32, f64),
    Fields{val: f64, check: Option<bool>, on: i8},
}

Answer

Method

Define a __new__ constructor with __getnewargs__ and construct a secondary enum internally to manage the correct arguments for each variant.

#[derive(Clone, Debug, FromPyObject, IntoPyObject)]
enum MyEnumNewArgs {
    Single(i32, f64),
    Fields(f64, Option<bool>, i8)
}

#[pymethods]
impl MyEnum {
    #[new]
    fn new_py(args: MyEnumNewArgs) -> MyEnum {
        match args {
            MyEnumNewArgs::Single(a, b) => MyEnum::Single(a, b),
            MyEnumNewArgs::Fields(a, b, c) => MyEnum::Fields{val: a, check: b, on: c},
        }
    }

    fn __getnewargs__<'py>(&self, _py: Python<'py>) -> MyEnumNewArgs {
        match self {
            MyEnum::Single(a, b) => MyEnumNewArgs::Single(*a, *b),
            MyEnum::Fields{val: a, check: b, on: c} => MyEnumNewArgs::Fields(*a, b.clone(), *c),
        }
    }
}

This almost works.

But, due to the nested class structure of complex enums (for example(#5250 (reply in thread))) they seemed to exhibit similar issues with pickling NamedTuples, e.g. (https://stackoverflow.com/questions/4677012/python-cant-pickle-type-x-attribute-lookup-failed)

Attempting to now pickle these objects will result in the error:

_pickle.PicklingError: Can't pickle <class 'mymodule.rs.MyEnum_Single'>: attribute lookup MyEnum_Single on mymodule.rs failed

If the namespace for the module is monkey-patched at top level

import mymodule.rs
mymodule.rs.MyEnum_Single = mymodule.rs.MyEnum.Single
mymodule.rs.MyEnum_Fields = mymodule.rs.MyEnum.Fields

Then pickling will succeed.

single = MyEnum.Single(3, 4)
fields = MyEnum.Fields(3, None, 2)
import pickle
for obj in [single, fields]:
    p = pickle.dumps(obj)
    l = pickle.loads(p)
    assert obj == l  # WORKS: AT BLOODY LAST!

Gotchas

If any of your variants have only 1 arguments then IntoPyTuple doesnt seem to play nice, the conversion created the object rather than a 1-element tuple and Python complained. In these cases I had to manually implement IntoPyObject (v25), for example:

impl<'py> IntoPyObject<'py> for MyEnumNewArgs {
    type Target = PyTuple;
    type Output = Bound<'py, Self::Target>;
    type Error = std::convert::Infallible;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            MyEnumNewArgs::Single(a) => Ok((a,).into_pyobject(py).unwrap()),
            MyEnumNewArgs::Fields(a, b, c) => Ok((a,b,c).into_pyobject(py).unwrap()),
        }
    }
}

I also do not know if monkey-patching the namespace has any other unwanted consequences.

Failures

I have not figured out how to do this for empty variants. I suppose I might need to manually implement a FromPyObject for MyEnumNewArgs. E.g.

#[pyclass(module = "mymodule.rs", eq)]
#[derive(Clone, Debug, PartialEq)]
pub enum MyEnum {
    Null(),
    Fields{val: f64, check: Option<bool>, on: i8},
}

#[derive(Clone, Debug)]  // need to implement FromPyObject and IntoPyObject
enum MyEnumNewArgs {
    Null(),
    Fields(f64, Option<bool>, i8)
}

[attack68]

Ok at this point I'm just updating this as my own reference.

Update

MyEnum is...	PyO3 pickling solution...	Case
Simple	Use U8 hard coded identifiers as described above	A
Complex and each variant has a unique set of inputs >= 2	For MyEnumNewArgs, FromPyObject and IntoPyObject can be derived. Namespace must be patched.	B
Complex and each variant has a unique set of inputs >= 1	For MyEnumNewArgs, FromPyObject can be derived and IntoPyObject must be overloaded to return a PyTuple. Namespace must be patched.	C
Complex and each variant has a unique set of inputs >= 0	For MyEnumNewArgs, FromPyObject and IntoPyObject must be overloaded. Namespace must be patched.	D
Complex and some variants have identical arg signature	Must add a u8 bytes identifier and proceed as with the above and the simple type	E

For the complex enum the pyclass is frozen meaning the state setting does not acquire the mutable reference. See error above,

A)

Simple Enum

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

#[pymethods]
impl MyEnum2 {
    #[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,))
    }
}

B)

Complex and each variant has a unique set of inputs >= 2

#[pyclass(eq)]
#[derive(Clone, Debug, PartialEq)]
enum MyEnum {
    A(u32, f64),
    B(u32, u32),
}

#[derive(FromPyObject, IntoPyObject)]
enum MyEnumNewArgs {
    A(u32, f64),
    B(u32, u32),
}

#[pymethods]
impl MyEnum {
    fn __getnewargs__(&self) -> MyEnumNewArgs {
        match self {
            MyEnum::A(x, y) => MyEnumNewArgs::A(*x, *y),
            MyEnum::B(x, y) => MyEnumNewArgs::B(*x, *y),
        }
    }

    #[new]
    fn new_py(args: MyEnumNewArgs) -> MyEnum {
        match args {
            MyEnumNewArgs::A(x, y) => MyEnum::A(x, y),
            MyEnumNewArgs::B(x, y) => MyEnum::B(x, y),
        }
    }
}

mymodule.rs.MyEnum_A = mymodule.rs.MyEnum.A
mymodule.rs.MyEnum_B = mymodule.rs.MyEnum.B

C)

Complex and each variant has a unique set of inputs >= 1

#[pyclass(eq)]
#[derive(Clone, Debug, PartialEq)]
enum MyEnum {
    A(u32),
    B(u32, u32),
}

#[derive(FromPyObject)]
enum MyEnumNewArgs {
    A(u32),
    B(u32, u32),
}

impl<'py> IntoPyObject<'py> for MyEnumNewArgs {
    type Target = PyTuple;
    type Output = Bound<'py, Self::Target>;
    type Error = std::convert::Infallible;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            MyEnumNewArgs::A(x) => Ok((x,).into_pyobject(py).unwrap()),
            MyEnumNewArgs::B(x, y) => Ok((x,y).into_pyobject(py).unwrap()),
        }
    }
}

#[pymethods]
impl MyEnum {
    fn __getnewargs__(&self) -> MyEnumNewArgs {
        match self {
            MyEnum::A(x) => MyEnumNewArgs::A(*x),
            MyEnum::B(x, y) => MyEnumNewArgs::B(*x, *y),
        }
    }

    #[new]
    fn new_py(args: MyEnumNewArgs) -> MyEnum {
        match args {
            MyEnumNewArgs::A(x) => MyEnum::A(x),
            MyEnumNewArgs::B(x, y) => MyEnum::B(x, y),
        }
    }
}

mymodule.rs.MyEnum_A = mymodule.rs.MyEnum.A
mymodule.rs.MyEnum_B = mymodule.rs.MyEnum.B

D)

Complex and each variant has a unique set of inputs >= 0

#[pyclass(module= "rateslib.rs", eq)]
#[derive(Clone, Debug, PartialEq)]
enum MyEnum {
    A(u32),
    B(),
}

enum MyEnumNewArgs {
    A(u32),
    B(),
}

impl<'py> IntoPyObject<'py> for MyEnumNewArgs {
    type Target = PyTuple;
    type Output = Bound<'py, Self::Target>;
    type Error = std::convert::Infallible;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            MyEnumNewArgs::A(x) => Ok((x,).into_pyobject(py).unwrap()),
            MyEnumNewArgs::B() => Ok(PyTuple::empty(py)),
        }
    }
}

impl<'py> FromPyObject<'py> for MyEnumNewArgs {
    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
        let ext: PyResult<(u32,)> = ob.extract();
        match ext {
            Ok(v) => Ok(MyEnumNewArgs::A(v.0)),
            Err(_) => Ok(MyEnumNewArgs::B()),
        }
    }
}

#[pymethods]
impl MyEnum {
    fn __getnewargs__(&self) -> MyEnumNewArgs {
        match self {
            MyEnum::A(x) => MyEnumNewArgs::A(*x),
            MyEnum::B() => MyEnumNewArgs::B(),
        }
    }

    #[new]
    fn new_py(args: MyEnumNewArgs) -> MyEnum {
        match args {
            MyEnumNewArgs::A(x) => MyEnum::A(x),
            MyEnumNewArgs::B() => MyEnum::B(),
        }
    }
}

mymodule.rs.MyEnum_A = mymodule.rs.MyEnum.A
mymodule.rs.MyEnum_B = mymodule.rs.MyEnum.B

E)

When a complex enum contains variants with the same arg signatures the FromPyObject cannot directly know which of these variants it will be. Therefore one additional argument is used as a type identifier. The final u8 value of each variant is used specifically as this identifier when pickling and for no other reason.

Using a fields type variant allows Python users to effectively ignore the item, although it makes it messier for Rust code.

#[pyclass(eq)]
#[derive(Clone, Debug, PartialEq)]
enum MyEnum {
    A(u32, u8),
    #[pyo3(constructor = (_pickle_id=1))]
    B{ _pickle_id: u8},
    C(u8),
}

enum MyEnumNewArgs {
    A(u32, u8),
    B(u8),
    C(u8),
}

impl<'py> IntoPyObject<'py> for MyEnumNewArgs {
    type Target = PyTuple;
    type Output = Bound<'py, Self::Target>;
    type Error = std::convert::Infallible;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            MyEnumNewArgs::A(x, _) => Ok((x,0).into_pyobject(py).unwrap()),
            MyEnumNewArgs::B(_) => Ok((1,).into_pyobject(py).unwrap()),
            MyEnumNewArgs::C(_) => Ok((2,).into_pyobject(py).unwrap()),
        }
    }
}

impl<'py> FromPyObject<'py> for MyEnumNewArgs {
    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
        let ext: PyResult<(u32, u8)> = ob.extract();
        if ext.is_ok() {
            return Ok(MyEnumNewArgs::A(ext.unwrap().0, 0))
        }
        let ext: (u8,) = ob.extract()?;
        match ext {
            (1,) => Ok(MyEnumNewArgs::B(1)),
            (2,) => Ok(MyEnumNewArgs::C(2)),
            _ => panic!("Undefined"),
        }
    }
}

#[pymethods]
impl MyEnum {
    fn __getnewargs__(&self) -> MyEnumNewArgs {
        match self {
            MyEnum::A(x, _) => MyEnumNewArgs::A(*x, 0),
            MyEnum::B{ _pickle_id: _ } => MyEnumNewArgs::B(1),
            MyEnum::C(_) => MyEnumNewArgs::C(2),
        }
    }

    #[new]
    fn new_py(args: MyEnumNewArgs) -> MyEnum {
        match args {
            MyEnumNewArgs::A(x, _) => MyEnum::A(x, 0),
            MyEnumNewArgs::B(_) => MyEnum::B{ _pickle_id: 1 },
            MyEnumNewArgs::C(_) => MyEnum::C(2),
        }
    }
}

mymodule.rs.MyEnum_A = mymodule.rs.MyEnum.A
mymodule.rs.MyEnum_B = mymodule.rs.MyEnum.B
mymodule.rs.MyEnum_C = mymodule.rs.MyEnum.C