Solved: slice protocol (and custom conversions) for a Rust object exposed to Python via pyo3

[chetmurthy]

I had posted earlier today about having trouble figuring out how to
implement the slice protocol for a Rust object exposed to Python via
pyo3. I finally figured it out, and for the sake of other newbies,
I thought I should write up the solution as well as the methodology
for figuring it out.

Because I'm sure others are going to want to solve similar problems.

For reference purposes, here's a little cargo project where I
developed the solution: rust-sandbox/lib.rs at main · chetmurthy/rust-sandbox · GitHub

The Problem

I have a Rust struct

#[pyclass]
#[repr(transparent)]
#[derive(Clone)]
pub struct Thing {
    members : Vec<(String, i64)>,
}

and I want to implement the Python slice protocol on it. That means I want to support
v[i] and v[start:stop:step]: the former returns (String, i64) and the latter a
vector of those.

With some searching I learned that pyo3 would convert the arguments to
__getitem__ into a Rust enum

#[derive(FromPyObject)]
enum SliceOrInt<'a> {
    Slice(&'a PySlice),
    Int(isize),
}

but it didn't help me with constructive the return-value. I could write a function

fn __getitem__(&self, idx: SliceOrInt, py: Python) -> PyResult<(String, i64)>

and that worked great, but that isn't enough: I needed to return that, or
PyResult<Vec<(String, i64)>>. And so I thought: I need to learn how to convert
each of these two types into PyObject -- then I can just return PyObject.

But everything I tried .... failed.

cargo expand is your friend

But then I found cargo expand, which, when you run it in your project, will macro-expand
your source-files and display them. So I learned that the way that (e.g.) (String, i64)
gets converted into PyObject is via the code

_pyo3::callback::convert(
	_py,
	Thing::__getitem__(_slf, arg0, _py),
)

And then digging into pyo3::callback::convert, I found that it used the IntoPy trait
in order to effect the conversion. So I first figured out how to write my own conversion
for (String, i64). But the code in this area is all happening with the Python "gil"
held, so it was clear I needed to write something that would get called from that
convert function. Then it hit me: I could write my own custom enum and define the
IntoPy trait on it.

enum SliceResult<T> {
    It(T),
    Slice(Vec<T>)
}
impl<T : IntoPy<PyObject>> IntoPy<PyObject> for SliceResult<T> {
    fn into_py(self, py: Python<'_>) -> PyObject {
        match self {
            SliceResult::It(it) => it.into_py(py),
            SliceResult::Slice(v) => {
                v.into_py(py)
            }
        }
    }
}

and then __getitem__ becomes

    fn __getitem__(&self, idx: SliceOrInt, py: Python) -> PyResult<SliceResult<(String, i64)>> {
        match idx {
            SliceOrInt::Slice(slice) => {
                let psi = slice.indices(self.members.len() as i64)? ;
                let (start, stop, step) = (psi.start, psi.stop, psi.step) ;
                let m : Vec<(String, i64)> =
                    self.members[start as usize..stop as usize].iter()
                    .step_by(step as usize)
                    .map(|p| (p.0.clone(), p.1))
                    .collect() ;
                let m = SliceResult::Slice(m) ;
                Ok(m)
            },
            SliceOrInt::Int(idx) => {
                (0 <= idx && idx < self.members.len() as isize).then(|| ())
                    .ok_or(PyException::new_err(format!("__getitem__ called on invalid index {}", idx))) ? ;
                let m = &self.members[idx as usize] ;
                let m = SliceResult::It((m.0.clone(), m.1)) ;
                Ok(m)
            }
        }
    }

What does that do?

1. match on the SliceOrInt; in the case of an Int, get the idx-th member, wrap it in a SliceResult::It and we're done.

2. in the case of a Slice, we need to extract the slice indices
   (searching for the type PySlice led to the type PySliceIndices
   which holds the relevant start/stop/step, and that you had to pass
   in the length of your vector, so that it could adjust the indices
   for you to make them legal. Which nice. Then just iterate over
   your vector, collect the results, and wrap in a SliceResult::Slice.

And because you've implemented IntPy on your custom enum, the
conversion machinery will take care of converting it to Python
types.

In the end, it wasn't very hard. What I needed, was to look at the
code of pyo3::callback::convert and figure out what it did.