Add partial generic support in pyfunctions

[Bben01]

Hi,

I am currently developing a crate that export a generic function that is meant to be configured by users at compile time before being exported to python. This looks like this:

trait Config {
    const KEY: [u8; 32];
}

pub fn encrypt<C: Config>(data: &[u8]) -> Vec<u8> {
    // use C::KEY here to encrypt the data  
}

Then the user of that crate would export that function in their python module using wrap_pyfunction! macro.

Currently, this is not possible. I have to export a declarative macro that will construct that function in the user's module, so that it can avoid the generic and use C directly.

Would you accept a PR that will modify the pyfunction proc macro and the wrap_pyfunction! declarative macro to support generics, where the user has to define the name + all the generics of the function in the wrap_pyfunction macro if the function is generic?

This would look like that:

trait Config {
    const KEY: [u8; 32];
}

#[pyfunction]
pub fn encrypt<C: Config>(data: &[u8]) -> Vec<u8> {
    // use C::KEY here to encrypt the data  
}

// -- in user crate --

struct Conf;

impl Config for Conf {
    const KEY: [u8; 32] = [0; 32];
}

#[pymodule]
fn module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!("encrypt", encrypt, Conf)?)
}

(The exact syntax can be changed)

[mejrs]

Would you accept a PR...

Before a PR I'd like to see a sketch of what the design might look like. What does it look like from the Python side, what do the macros do, how does it compose when there are multiple generics, that kind of stuff. So for example, if I do...

#[pymodule]
fn module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!("encrypt", encrypt, Conf)?);
    m.add_function(wrap_pyfunction!("encrypt", encrypt, OtherConf)?)
}

then what happens? Do I get a module with functions encrypt_conf and encrypt_otherconf? Do I get an encrypt function that checks the type and forwards either to the Conf or OtherConf monomorphization?

While there are some tricks, the #[pyfunction] macro and wrap_pyfunction! macros can't look at each other and agree how they're dealing with a generic function.

[mejrs]

Note also that this isn't really what people will think of with "I'd like to expose my generic function to Python".

They'll imagine a function with generic arguments, like

fn blah<T>(x: T){}

So any proposal should cover this case as well.

[Bben01]

I will give a more detailed explanation tomorrow, but for your question about the example you gave, it will result in a compile error, since you are trying to define the same name twice (the first argument to the wrap_pyfunction is the name that will be visible from python).

About the multiple generics, you will be able to define multiple generics (and specify them all in the wrap macro)

[mejrs]

As far as the use case "i'd like to have my python function accept multiple types" goes, that is semi-ergonomic by using https://pyo3.rs/v0.24.1/conversions/traits.html#deriving-frompyobject-for-enums. I'm also open for improvements to that, if the design is elegant enough.

[Bben01]

Ok, here is the more detailed version:

The idea behind that feature is that the user of my crate wants to expose a python function that is NOT generic, but "configured" in the rust side, so that the python side doesn't even know the configuration

It could be expanded to support more cases.

From the python side, you don't see anything. From the macro side, we modify the trampoline to be generic over the same generics that the pyfunction defines, then when we wrap the pyfunction in the module, we specify all the generics. It does mean that we need to change the currently const _PYO3_DEF to a fn _PYO3_DEF, and the wrap_pyfunction will call that function instead of passing the const to WrapPyFunctionArg::wrap_pyfunction.

The user (that wrap the pyfunction) will have to give the function a name, so if it give the same name twice (like in your example), the second one will overwrite the first one, like if you wrapped the same pyfunction in one module.

We could also support generics in parameters and return type if the use makes sure that they implement FromPyObject/IntoPyObject (we could also add the bound automatically).

About the API, we can do it in multiple ways:

1. be explicit, require the use to write #[pyfunction(generic)] to allow it to be generic, and use a corresponding wrap_generic_pyfunction! macro
2. be implicit about it, write #[pyfunction] like a regular function, and modify the wrap_pyfunction! macro to support the generics

Currently, this doesn't cover cases where the function accept multiple types.

[mejrs]

It seems to me that you can already do that today:

pub fn define_decrypt<C: Config>(py: Python<'_>) -> PyResult<PyCFunction> {
    let f = |data: &[u8]| {
           decrypt::<C>(data)
    };
    PyCFunction::new_closure(py, Some(c"decrypt"), Some(c"decrypts something"),  f)
}

#[pymodule]
fn module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add("decrypt", define_config::<Conf>())
}

[Bben01]

I tried this:

use pyo3::prelude::*;
use pyo3::types::{PyCFunction, PyDict, PyTuple, PyTupleMethods};
use pyo3::Bound;

pub trait Config {
    const OFFSET: i32 = 0;
}

struct C;

impl Config for C {
    const OFFSET: i32 = 1;
}

/// A simple function that adds a number to a constant offset.
pub fn add<C: Config>(number: i32) -> i32 {
    number + C::OFFSET
}

pub fn wrap_pygeneric_function<C: Config>(py: Python) -> PyResult<Bound<PyCFunction>> {
    let f = |args: &Bound<'_, PyTuple>, _kwargs: Option<&Bound<'_, PyDict>>| {
        add::<C>(args.get_item(0).unwrap().extract().unwrap())
    };
    PyCFunction::new_closure(
        py,
        Some(c"add"),
        Some(c"A simple function that adds a number to a constant offset."),
        f,
    )
}

#[pymodule]
fn pyexample(m: &Bound<PyModule>) -> PyResult<()> {
    m.add("add", wrap_pygeneric_function::<C>(m.py())?)?;
    Ok(())
}

This is technically possible, but notice that I have to copy the doc comment and parse the arguments myself.

If pyo3 exported a public API to do argument parsing, I could write a small proc macro that takes the name, doc and arguments and generate wrap_pygeneric_function.

I still thinks that integrating this in pyo3 could make things easier.