Provide Python API for Rust code that uses ADTs and trait objects

[cdiwoxv]

Hi PyO3 community,

First of all, huge thanks to everyone that makes PyO3 happen.

In our codebase we have a Rust library that we'd like to provide Python APIs for other internal users to use. A simplified version of the public interface of the library is something like this:

pub struct EventA {
    pub field_a1: i64,
    pub field_a2: bool,
}

pub struct EventB {
    pub field_b1: f64,
    pub field_b2: i32,
}

pub enum Event {
    EventA(EventA),
    EventB(EventB),
}

pub trait Client {
    fn handle_event(&mut self, event: &Event);
}

pub struct Driver {
    clients: Vec<Box<dyn Client>>,
}

impl Driver {
    pub fn new(clients: Vec<Box<dyn Client>>) -> Self {
        Self { clients }
    }

    fn emit_event(&mut self, event: Event) {
        for client in self.clients.iter_mut() {
            client.handle_event(&event)
        }
    }

    pub fn emit_event_a(&mut self, field_a1: i64, field_a2: bool) {
        // complicated things happen in here in real codebase
        let event = Event::EventA(EventA { field_a1, field_a2 });
        self.emit_event(event)
    }

    pub fn emit_event_b(&mut self, field_b1: f64, field_b2: i32) {
        // complicated things happen in here in real codebase
        let event = Event::EventB(EventB { field_b1, field_b2 });
        self.emit_event(event)
    }

The way a user uses this library is, the user will have its own type that implements Client, then construct a Driver with its own clients, and then call emit_event_* functions.

I'd like to provide a pythonic interface for this library (so Client can be implemented in Python). The ideal Python interface in my mind would be something like this:

from abc import Abc, abstractmethod
from dataclasses import dataclass
from typing import List

class Event(Abc):
    pass

@dataclass
class EventA(Event):
    field_a1: int
    field_a2: bool

@dataclass
class EventB(Event):
    field_b1: float
    field_b2: int

class Client(Abc):
    @abstractmethod
    def handle_event(self, event: Event):
        pass

class Driver:
    def __init__(self, clients: List[Client]):
        pass # should construct a Rust driver
    def emit_event_a(self, field_a1: int, field_a2: bool):
        pass # should use the Rust driver
    def emit_event_b(self, field_b1: float, field_b2: int):
        pass # should use the Rust driver

Then a user would construct its own client and use the library this way:

class MyClient(Client):
    def handle_event(self, event: Event):
        match event:
            case EventA(field_a1=a1, field_a2=a2): print(f"got A(a1={a1}, a2={a2})")
            case EventB(field_b1=b1, field_b2=b2): print(f"got B(b1={b1}, b2={b2})")

driver = Driver([MyClient()])
driver.emit_event_a(100, False)
driver.emit_event_b(3.4, 20)

I have been reading PyO3 docs and playing with it for a few days now. I have two main questions:

1. How do I generate the classes Event, EventA, EventB with PyO3?

I've tried something like this, first create a dummy base class:

#[pyclass(name="Event", subclass)]
struct PyEvent {}

Then I can add #[pyclass(extends=PyEvent)] to EventA and EventB.

I thought I could then implement IntoPy<PyObject> for Event but calling EventA::into_py and EventB::into_py. But I found that they no longer work because the trait bounds AsPyPointer was not satisfied.

I can see a compromising way is to define

Event = typing.Union[EventA, EventB]

in Python instead of using base class.

But I wonder whether there is some way to make the desired class hierarchy in Python side work?

2. Separately, I would like a Client to be implemented in Python side, and use the Rust Driver to send events to Python side clients. The concepts of an abstract base class in Python and a trait in Rust seem analogous enough, but I failed to come up with code that uses PyO3 to achieve what I had in mind. Can someone point me to a right direction?

Constraints:

• It's possible for us to change the interface of the Rust library without sacrificing its semantics or performance, and without sacrificing the ergonomics at Rust side too much.
• We use Rust stable and Python 3.10, Linux only.

I'll really appreciate if my question gets answered. A repo with my sample code is in here: https://github.com/cdiwoxv/pyo3-question-2535

[messense]

I thought I could then implement IntoPy for Event but calling EventA::into_py and EventB::into_py. But I found that they no longer work because the trait bounds AsPyPointer was not satisfied.

All types in PyO3 implement IntoPy, as does a #[pyclass] which doesn't use extends, see also #1836

[cdiwoxv]

Thank you @messense . I understand 1 is due to current limitation of PyO3. I currently work around it in my code by using

Event = typing.Union[EventA, EventB]

in Python side. Which isn't the same as a base class but seems good enough for my use case for now.

Can you please shed some light on my question 2? Thank you!

[messense]

For question 2, I've come up something like this, it compiles but I'm not sure if it's a good solution (full code):

#[pyclass(name = "Client", subclass)]
#[derive(Clone)]
struct PyClient {}

#[pymethods]
impl PyClient {
    #[new]
    fn new() -> Self {
        Self {}
    }

    fn handle_event(&mut self, _event: &PyAny) -> PyResult<()> {
        Err(PyNotImplementedError::new_err("not implemented"))
    }
}

impl Client for PyClient {
    fn handle_event(&mut self, event: &Event) {
        Python::with_gil(|py| {
            let py_event = event.to_object(py);
            // TODO: error handling
            self.handle_event(py_event.as_ref(py)).unwrap();
        });
    }
}

#[pyclass(name = "Driver", unsendable)]
struct PyDriver {
    inner: Driver,
}

#[pymethods]
impl PyDriver {
    #[new]
    fn new(py: Python<'_>, clients: Vec<PyClient>) -> PyResult<Self> {
        let clients: Vec<_> = clients
            .into_iter()
            .map(|c| Box::new(c) as Box<dyn Client>)
            .collect();
        Ok(Self {
            inner: Driver::new(clients),
        })
    }
}

User needs to subclass Client on Python side.

[cdiwoxv]

Unfortunately, I don't think this solution works.

I updated my repo with a test that fails: cdiwoxv/pyo3-question-2535@3c2050f

In particular, I guess that when PyDriver::new is called, PyClient.clone() is called. And the vec<Box<dyn Client>> we put into Driver is simply a bunch of empty PyClient {}s. So calling handle_event will raise "not implemented".

[messense]

Right, I was also confused when missing #[derive(Clone)] on PyClient https://github.com/cdiwoxv/pyo3-question-2535/blob/3c2050fd26e50f6bed9b195decb799faddaa23d2/src/lib.rs#L85 , the error message was quite bad:

error[E0277]: the trait bound `Vec<PyClient>: FromPyObject<'_>` is not satisfied
   --> src/lib.rs:118:38
    |
118 |     fn new(_py: Python<'_>, clients: Vec<PyClient>) -> PyResult<Self> {
    |                                      ^^^ the trait `FromPyObject<'_>` is not implemented for `Vec<PyClient>`
    |
    = help: the trait `FromPyObject<'a>` is implemented for `Vec<T>`
note: required by a bound in `extract_argument`
   --> /Users/messense/.cargo/registry/src/github.com-1ecc6299db9ec823/pyo3-0.16.5/src/impl_/extract_argument.rs:14:8
    |
14  |     T: FromPyObject<'py>,
    |        ^^^^^^^^^^^^^^^^^ required by this bound in `extract_argument`


error: aborting due to previous error==> ] 41/42: pyo3-question


For more information about this error, try `rustc --explain E0277`.

error: could not compile `pyo3-question` due to 2 previous errors

[davidhewitt]

You probably want to take Vec<Py<PyClient>> as this will be the Python objects (including all their subclass information).

Similarly you'll need to impl Client for Py<PyClient> and call the Python handle_event implementation (py_client.call_method1((event,))) so that the Python subclass overrides are dispatched to correctly.