How can I implement python operators in PyO3

Question

I'm trying to implement a vector class in rust for my math library.

#[pyclass]
struct Vec2d {
    #[pyo3(get, set)]
    x: f64,
    #[pyo3(get, set)]
    y: f64
}

But I can't figure out how I can overload the standard operators (+, -, *, /)

I Tried implementing the Add trait from std::ops with no luck

impl Add for Vec2d {
    type Output = Vec2d;
    fn add(self, other: Vec2d) -> Vec2d {
        Vec2d{x: self.x + other.x, y: self.y + other.y }
    }
}

I also tried adding __add__ method to the #[pymethods] block

fn __add__(&self, other: & Vec2d) -> PyResult<Vec2d> {
    Ok(Vec2d{x: self.x + other.x, y: self.y + other.y })
}

but still does not work.

With the second approach I can see that the method is there, but python doesn't recognize it as operator overload

In [2]: v1 = Vec2d(3, 4)
In [3]: v2 = Vec2d(6, 7)
In [4]: v1 + v2
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-4-08104d7e1232> in <module>()
----> 1 v1 + v2

TypeError: unsupported operand type(s) for +: 'Vec2d' and 'Vec2d'

In [5]: v1.__add__(v2)
Out[5]: <Vec2d object at 0x0000026B74C2B6F0>

Abdul Niyas P M · Accepted Answer · 2020-09-03T04:08:02.303

As per the PyO3 Documentation,

Python's object model defines several protocols for different object behavior, like sequence, mapping or number protocols. PyO3 defines separate traits for each of them. To provide specific python object behavior you need to implement the specific trait for your struct.

Important note, each protocol implementation block has to be annotated with #[pyproto] attribute.

__add__, __sub__ etc are defined within PyNumberProtocol Trait.

So you could implement PyNumberProtocol for your Vec2d struct to overload standard operations.

#[pyproto]
impl PyNumberProtocol for Vec2d {
    fn __add__(&self, other: & Vec2d) -> PyResult<Vec2d> {
            Ok(Vec2d{x: self.x + other.x, y: self.y + other.y })
   }
}

This solution is not tested, For the complete working solution check @Neven V's answer.

Neven V. · Answer 2 · 2020-06-07T19:13:00.863

I'll add this answer to spare others from searching for hours like I did.

Using the answer provided by @Abdul Niyas P M, I had the following error:

error: custom attribute panicked
  --> src/vec2.rs:49:1
   |
49 | #[pyproto]
   | ^^^^^^^^^^
   |
   = help: message: fn arg type is not supported

It turns out that this cryptic error message hides two problems. The first problem is that __add__ should take values and not references, thus we remove the & before self and Vec2. This allows us to get rid of the error message:

error[E0277]: the trait bound `&vec2::Vec2: pyo3::pyclass::PyClass` is not satisfied
   --> src/vec2.rs:47:1
    |
47  | #[pyproto]
    | ^^^^^^^^^^ the trait `pyo3::pyclass::PyClass` is not implemented for `&vec2::Vec2`
    |

The second of these problems can be revealed when we specify the type of self:

// DOES NOT COMPILE
#[pyproto]
impl PyNumberProtocol for Vec2 {
    fn __add__(self: Vec2, other: Vec2) -> PyResult<Vec2> {
        Ok(Vec2{x: self.x + other.x, y: self.y + other.y})
    }
}

This fails to compile with the errror message

error[E0185]: method `__add__` has a `self: <external::vec3::Vec3 as pyo3::class::number::PyNumberAddProtocol<'p>>::Left` declaration in the impl, but not in the trait
  --> src/external/vec2.rs:49:5
   |
49 |     fn __add__(self: Vec2, other: Vec2) -> PyResult<Vec2> {
   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `self: <vec2::Vec2 as pyo3::class::number::PyNumberAddProtocol<'p>>::Left` used in impl
   |
   = note: `__add__` from trait: `fn(<Self as pyo3::class::number::PyNumberAddProtocol<'p>>::Left, <Self as pyo3::class::number::PyNumberAddProtocol<'p>>::Right) -> <Self as pyo3::class::number::PyNumberAddProtocol<'p>>::Result`

This leads us to the final working solution (as of June 2020):

#[pyproto]
impl PyNumberProtocol for Vec2 {
    fn __add__(lhs: Vec2, rhs: Vec2) -> PyResult<Vec2> {
        Ok(Vec2{x: lhs.x + rhs.x, y: lhs.y + rhs.y})
    }
}

Which successfully compiles under Rust nightly 1.45 and has been checked to work from Python.

It is also possible to have rhs of another type:

#[pyproto]
impl PyNumberProtocol for Vec2 {
    fn __mul__(lhs: Vec2, rhs: f64) -> PyResult<Vec2> {
        Ok(Vec3{x: lhs.x * rhs, y: lhs.y * rhs})
    }
}

Also note that having self is not always a problem:

#[pyproto]
impl PyNumberProtocol for Vec2 {
    fn __neg__(self) -> PyResult<Vec2> {
        Ok(Vec3{x: -self.x, y: -self.y})
    }
}

How can I implement python operators in PyO3

2 Answers2