Dynamic dispatch with Python objects

[diegogangl]

Hi, I'm trying to pass custom objects through Python to Rust. I need to store one of these in a main struct. The Python classes have some properties of their own and a method I'm calling from the main struct.

This is the Python API I'm shooting for:

main = rust.MainClass()
sub_1 = rust.SubClass01()
sub_2 = rust.SubClass02()

sub_1.a_value = 2
sub_2.a_boolean = True

main.inner_class = sub_1
print(main.do_the_thing())

This is the code I have on the rust side.

#[pyclass]
#[derive(Clone)]
pub struct BaseClass {}

#[pymethods]
impl BaseClass {
    fn do_something(&self) -> f64 {
        -1.0
    }

    #[new]
    fn new() -> Self {
        BaseClass {}
    }

}


#[pyclass(extends=BaseClass)]
#[derive(Clone)]
pub struct SubClass01 {

    #[pyo3(get, set)]
    a_value: u8,
}

#[pymethods]
impl SubClass01 {
    fn do_something(&self) -> f64 {
        1.0
    }

    #[new]
    fn new() -> (Self, BaseClass) {
        (SubClass01 { a_value: 0 }, BaseClass::new())
    }
}


#[pyclass(extends=BaseClass)]
#[derive(Clone)]
pub struct SubClass02 {

    #[pyo3(get, set)]
    a_boolean: bool,
}

#[pymethods]
impl SubClass02 {
    fn do_something(&self) -> f64 {
        if self.a_boolean {
            2.0
        } else {
            0.0
        }
    }

    #[new]
    fn new() -> (Self, BaseClass) {
        (SubClass02 { a_boolean: true }, BaseClass::new())
    }
}


#[pyclass]
pub struct MainClass {

    #[pyo3(get, set)]
    inner_class: BaseClass
}

#[pymethods]
impl MainClass {

    fn do_the_thing(&self) -> f64 {
        self.inner_class.do_something()
    }

    #[new]
    fn new() -> Self {
        MainClass { inner_class: BaseClass {} }
    }
}

The problem is that the subclasses are never set. MainClass always calls BaseClass (the print at the end prints -1.0). I can sort of see why this happens on the rust side since I'm setting BaseClass by default and as the type (I'm guessing the subclasses are getting downcasted?).

I trait objects would be they way to solve this but I don't see how that works in the Pyo3/python side.

Here's what I tried to do with trait objects:

// Yeah this is a terrible trait name :)
pub trait Instanceable {
    fn do_some(&self) {}
}

impl Instanceable for SubClass01 {
    fn do_some(&self) { println!("555"); }
}

impl Default for SubClass01 {
    fn default() -> Self {
        SubClass01 { a_value: 0 }
    }
}

#[pyclass]
pub struct MainClass {
    some_class: Box<dyn Instanceable>
}


#[pymethods]
impl MainClass {

    #[new]
    fn new() -> Self {
        MainClass {
            some_class: Box::new(SubClass01::default())
        }
    }

    #[setter(some_class)]
    fn set_class<T: Instanceable>(&mut self, cls: T) -> PyResult<()> {
        Box::new(cls);
        Ok(())
    }
}

In this case I can't compile because "a python method can't have a generic type parameter", which makes sense.
How do I pass a custom object to the setter? Is this kind of API or dynamic dispatch possible with PyO3?

Thanks!

[davidhewitt]

Hi @diegogangl, thanks for this, it's a good question.

You are right that Rust's mechanism of doing this kind of polymorphism is through trait objects. If you're looking to keep Python's semblance of shared mutability, you're probably looking for Rc<dyn Instanceable>.

With that in mind, I have a couple of ideas which should help you move forward:

Option 1: Trait Objects

As Rc<dyn Instanceable> isn't able to be a #[pyclass] natively, we need to wrap it up in one. In the example below, I've called it Instantiator.

The downside of this approach is on the Python side, where you have to add .instantiator() methods to each subclass, and you can't tell which subclass is currently in MainClass.

use std::rc::Rc;
use std::cell::RefCell;

pub trait Instanceable {
    fn do_some(&self) {}
}

#[pyclass]
#[derive(Clone)]
pub struct Instantiator {
    obj: Rc<dyn Instanceable>
}


struct SubClass01Inner {
    a_value: RefCell<u8>,
}

impl Instanceable for SubClass01Inner {
    fn do_some(&self) { println!("{}", self.a_value.borrow()); }
}

#[pyclass]
#[derive(Clone)]
pub struct SubClass01 {
    inner: Rc<SubClass01Inner>
}

#[pymethods]
impl SubClass01 {
    #[new]
    fn new() -> Self {
        SubClass01 {
            inner: Rc::new(SubClass01Inner { a_value: RefCell::new(0) })
        }
    }

    #[getter(a_value)]
    fn get_a_value(&self) -> u8 {
        *self.inner.a_value.borrow()
    }

    #[setter(a_value)]
    fn set_a_value(&mut self, value: u8) {
        self.inner.a_value.replace(value);
    }

    fn instantiator(&self) -> Instantiator {
        Instantiator {
            obj: self.inner.clone()
        }
    }
}

impl Default for SubClass01 {
    fn default() -> Self {
        SubClass01::new()
    }
}

#[pyclass]
pub struct MainClass {
    #[pyo3(get, set)]
    instantiator: Instantiator
}


#[pymethods]
impl MainClass {

    #[new]
    fn new() -> Self {
        MainClass {
            instantiator: SubClass01::default().instantiator()
        }
    }

    fn do_the_thing(&self) {
        self.instantiator.obj.do_some()
    }
}

This would give you an experience in Python like so:

>>> import rust
>>> m = rust.MainClass()
>>> m.do_the_thing()
0
>>> s = rust.SubClass01()
>>> m.instantiator = s.instantiator()
>>> s.a_value = 5
>>> m.do_the_thing()
5

Option 2: Python Dispatch

The alternative to the above is to give up on the dispatch on the Rust side, and instead resolve the right subclass method to call using a Python method call.

To do that, you just need to modify the MainClass in your first example to look like this:

#[pyclass]
pub struct MainClass {

    #[pyo3(get, set)]
    inner_class: PyObject
    // The above could potentially be Py<BaseClass> to get a little type safety, but there's a
    // missing trait implementation which I'll resolve shortly...
}

#[pymethods]
impl MainClass {

    fn do_the_thing(&self, py: Python) -> PyResult<PyObject> {
        self.inner_class.getattr(py, "do_something")?.call0(py)
    }

    #[new]
    fn new(py: Python) -> Self {
        MainClass { inner_class: BaseClass {}.into_py(py) }
    }
}

With that, you get a Python API that looks a bit more natural to those users:

>>> import pyo3_scratch
>>> m = pyo3_scratch.MainClass()
>>> m.do_the_thing()
-1.0
>>> s = pyo3_scratch.SubClass01()
>>> m.inner_class = s
>>> m.do_the_thing()
1.0

[davidhewitt]

However, I think that neither solution is ideal - the trait object route is very messy at the moment, and the python dispatch route gives up most of the typing on the Rust side.

I'll let you know if I have any further thoughts to refine this with pyo3's existing implementation.

This kind of polymorphic dispatch is also a pattern that does crop up in Python APIs from time-to-time, so it's probably something we want to think about supporting better in pyo3. Marking as help-wanted in case anyone else has ideas what a nice solution for this could look like.

[diegogangl]

Hey @davidhewitt , thanks for looking into this!

Option 1 would the best, since I'd rather not give up Rust typing. What worries me about the 1st option is how memory management seems to be getting more complicated with borrows, derefs, etc.
I'm still a bit of a Rust noob TBH and I wonder what unforeseen issues I'd get if I go that route.

I think another option would be handling all the dispatch from Rust's side. Like having a method in Mainclass that takes a dict with the name of the class and it's properties, and then initializes the correct class. I wouldn't be able to call the methods from outside, but that's something I can live without for now (and I could make an extra pyclass to wrap around the struct if needed).

Another sub-option could be making a config struct (that is a pyclass). Then I can build that in Python, pass it to the setter in MainClass and the setter would take care of initializing the appropriate struct (based on the config struct type) and passing the config to it. Then in the "subclass" I can call that config like self.config.some_value. That would be a little more strict than the dict.

[davidhewitt]

Yep, I only showed one possible way that option 1 could be arranged - you could tweak that kind of thing as suits you to get an API more like what you're happy with.

Would be interested to see what your final solution looks like - maybe we can use it as experience to improve this part of pyo3.

[diegogangl]

Sure, I will post back here once I have some final code 👍

[ChrisPattison]

Are there any further thoughts on this feature? This sort of pattern is a very natural API to write from the Rust perspective

[davidhewitt]

I haven't had time to think about it myself. All input welcome!

[anand-bala]

I've been working on something that may need this feature soon and had a weird idea/question.

Would it be possible to define a macro (say #[pyprotocol]) that does structural subtyping like how a runtime_checkable Protocol does? Specifically, the macro would define a new runtime_checkable Protocol (not sure how to do this with PyO3) that matches the trait along with a trampoline class that implements the said trait and is constructed from a PyObject that implements the protocol (using isinstance on the Rust side). Then, the MainClass struct would just have to hold a trait object.

This may be wishful thinking, but this could potentially also allow a user to define some custom class on the Python side that implements said protocol and use it with a Rust implementation that uses trait objects.

EDIT: I will try sending a minimal example code as soon as I flesh this out in my head 😅 .

[alexcwyu]

Any update to support Trait Objects?