Method overloading with Rust traits

[jblindsay]

Hello, I have a Rust struct that I have exposed to Python using pyo3. I would like to implement overloaded methods using the typical Rust traits-based scheme, to allow a common named method to accept multiple types, e.g. add(f64), add(f32) but more complex than this. My problem is that pyo3 doesn't seem to allow traits on pymethods. I get the error:

error: #[pymethods] cannot be used on trait impl blocks

Is there no way to get around this to implement overloaded methods, or am I missing something? Thanks for your help.

[davidhewitt]

Can you share a minimal snippet of the code that you intuitively wrote which hit the error? I'd be happy to comment on the technical limitations, what you can do instead, and how we might be able to improve PyO3 in this area.

[jblindsay]

Thank you so much for your time. I have a struct called Raster, which, as the name implies, is effectively a large raster data structure like a 2D array. I'd like to be able to perform various ops on them, like add, subtract, multiply, etc. two Rasters together, or a Raster and a single f64, etc. Here's an example that actually caused the above stated error when used:

#[pymethods]
impl ops::Add<f64> for Raster {
    type Output = Raster;

    fn add(self, rhs: f64) -> Raster {
        let func = |a, b| a + b;

        let mut configs = self.configs.clone();
        configs.nodata = -32768.0;
        configs.data_type = DataType::F32;
        configs.photometric_interp = PhotometricInterpretation::Continuous;
        let mut output = Raster::initialize_using_config("", &configs);
        
        self.two_parameter_function(rhs, func, &mut output);

        output
    }
}

Then, of course, I'd ideally like to be able to use this to, within Python, say something like:

new_raster = raster1 + raster2

or

new_raster = raster1 + 50.0

I'm not sure if that helps you or not, but essentially, I would ideally like to use Rust's trait system to overload operators. In reading some other things on this forum, I get the sense that the complexity of doing so may not make it possible for pyo3 to have this feature. If that's the case, then I completely understand, but I just wanted to make sure that I wasn't missing something obvious.

[davidhewitt]

Right, I see. This idea of using #[pymethods] on std::ops::Add is interesting.

In principle we could support it (with limitations such as if you imported the trait under a different name like OpsAdd the macro would not recognise it). However, given the Add trait is generic (essentially overloadable), this would create a technical challenge which would likely make the PyO3 framework quite complicated. IMO for now there are higher priority items to tackle first.

The recommended way you can achieve this at the moment is to just define your Python __add__ as normal. You can make it handle multiple types using an enum input:

#[derive(FromPyObject)]
enum RasterOrF64 {
    Raster(Raster),
    F64(f64)
}

#[pymethods]
impl Raster {
    fn __add__(&self, other: RasterOrF64) -> Self {
        match other {
            RasterOrF64::Raster(r) => self.clone() + r,
            RasterOrF64::F64(f) => self.clone() + f,
        }
    }
}

[jblindsay]

This is great, thank you. I didn't realize that I could simply create a Rust method called __add__ and it would be recognized as the equivalent Python __add__ method. The enum trick is also very interesting. I wonder what that will look like on the Python side. Will the Raster/f64 need to be wrapped in a RasterOrF64 object before passing it to __add__? Anyhow, I think this answer gets me what I'm looking for and I truly appreciate you taking the time.

[jblindsay]

Just a quick follow up, this works beautifully on the Python side. I can use:

result = raster1 + 50.0

Or:

result = raster1 + raster2

And it's very natural. Every time I think think that I can't be any more impressed with pyo3, I'm proven wrong. It's truly a remarkable library and I'm very grateful.

[jblindsay]

By the way, am I missing something or are the comparison operators, __gt__, __lt__, etc., not implemented? I see all the other mathematical operators implemented here except the comparison operators. Thanks.

[davidhewitt]

At the moment these are all implemented by a single __richcmp__ method - briefly documented here.

There is #2089 to unbundle these into separate methods to match Python, just nobody has been able to get around to that yet (sorry)!

[jblindsay]

Okay, great, thanks for the quick reply.