PyO3 / PyO3/pyo3

ENH: boiler plate code for `PyPickle` to support Enums and multithreaded pools

Open
#4,465 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Rust
Stars
16.2k
Forks
1k
Avg merge
2d 6h
Merged PRs (30d)
66

Description

As context, I rarely use pickle directly in Python, but apparently I have been using it indirectly in the below code:

from multiprocessing import Pool
func = partial(other_func, **kwargs)
p = Pool(NUM_CPUS)
results = p.map(func, some_list)

When my project was pure Python this went undetected. When I started porting some of my classes to Rust with PyO3 all of my tests passed except ones involving the above lines. Why? Becuase my [pyclass] objects were not picklable, apparently.

Thus, for someone not at all au fait with pickling, this lead me on a rather blind hunt. However, I came across a very useful issue on this board, and implemented that code.

Later, my [pyclass] become more complex and fields used to construct them also require enum also labelled as [pyclass]. A simple example which works directly in Rust as MyEnum::Variant1 and Python as MyEnum.Variant1 is:

#[pyclass(module = "mymodule.rs")]
#[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum MyEnum {
    Variant1,
    Variant2,
    Variant3,
}

My method for pickling broke becuase simple enum do not have, or naturally need, a constructor method. To solve this I have used the following structure:

#[pymethods]
impl MyEnum {
    // Pickling
    #[new]
    fn new_py(variant: u8) -> PyResult<MyEnum> {
        match ad {
            0_u8 => Ok(MyEnum::Variant1),
            1_u8 => Ok(MyEnum::Variant2),
            2_u8 => Ok(MyEnum::Variant3),
            _ => Err(PyValueError::new_err("unreachable code on MyEnum pickle."))
        }
    }
    fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
        *self = deserialize(state.as_bytes()).unwrap();
        Ok(())
    }
    fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        Ok(PyBytes::new_bound(py, &serialize(&self).unwrap()))
    }
    fn __getnewargs__<'py>(&self) -> PyResult<(u8,)> {
        match self {
            MyEnum::Variant1 => Ok((0_u8,)),
            MyEnum::Variant2 => Ok((1_u8,)),
            MyEnum::Variant3 => Ok((2_u8,)),
        }
    }
}

I don't really want this code in my project. Its complicated (for me) and excessive when applied across all my [pyclass], and it is only needed for that pool multithreading. Issue raised to provoke any discussion if any of this can be abstracted away...

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing the PyO3 #[pyclass] and #[pymethods] mechanisms involved in the enum pickling example, then trace how Python multiprocessing.Pool requires objects to be picklable. Define what an abstraction would need to support for enums and multithreaded pools, and consider the example's successful round trip as the completion criterion.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, rust
Domain
api
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.