function defined by pyo3 is not pickable

[binh-vu]

Right now I can't pickle function created pyo3 due to the error: attribute lookup on <module_or_builtin> failed.

For pyo3 version 0.16.5, class created by pyo3 can pickle by changing #[pyclass] to #[pyclass(module = ...)] (this has not documented yet). However, the same solution does not work for #[pyfunction].

I have to manually set the __module__ attribute for it to work:

let func = wrap_pyfunction!(<func>, module)?;
func.setattr("__module__", "<module_path>")?;

It would be nice to have #[pyfunction(module =...)] working as pyo3 class.

[davidhewitt]

I just checked this myself and pickling appears to work fine.

Source code:

use pyo3::{prelude::*};

#[pyfunction]
fn foo() -> i32 {
    42
}

/// A Python module implemented in Rust.
#[pymodule]
fn pyo3_scratch(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(foo, m)?)?;
    Ok(())
}

Seems to work fine to me:

>>> import pickle
>>> import pyo3_scratch
>>> pickle.dumps(pyo3_scratch.foo)
b'\x80\x04\x95%\x00\x00\x00\x00\x00\x00\x00\x8c\x19pyo3_scratch.pyo3_scratch\x94\x8c\x03foo\x94\x93\x94.'

I will close this issue. If you are still observing a problem please reduce to a minimal repro and I can try to assist further.

[binh-vu]

Thank you for your reply. I forgot to mention that it happens to functions defined in submodule.

I create a repo that you can reproduce the issue here: pyo3-2469

This is the source code:

use pyo3::prelude::*;

/// square of a number
#[pyfunction]
fn square(x: f64) -> f64 {
    x * x
}

/// cube of a number
#[pyfunction]
fn cube(x: f64) -> f64 {
    x * x * x
}

pub fn register_submodule(py: Python<'_>, m: &PyModule) -> PyResult<()> {
    let submodule = PyModule::new(py, "sub1")?;

    submodule.add_function(wrap_pyfunction!(cube, submodule)?)?;
    m.add_submodule(submodule)?;

    Ok(())
}

#[pymodule]
fn demo(py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(square, m)?)?;
    register_submodule(py, m)?;
    Ok(())
}

Python:

import pickle
from demo import square, sub1

print(square(5))
print(square(10))
print(sub1.cube(10))

# error
pickle.dumps(sub1.cube)

# also error
cube = sub1.cube
pickle.dumps(cube)

# cannot import without using sys.modules trick
from demo.sub1 import cube

[davidhewitt]

Thanks - I think for now because it's a submodule and we already need to have tricks to make them work correctly as per #759, I'm going to leave a note in that issue and for now I'm afraid workarounds like you suggest in the OP are probably the right solution.

Once we have better submodule support this should hopefully not be an issue.