pyo3_segfault.rs

use pyo3::prelude::*;
use rayon::prelude::*;

struct PyStringWrapper {
    py_bytes: Py<PyString>,
    backed_str: PyBackedStr,
}

impl PyStringWrapper {
    fn from_bound(bound: Bound<'_, PyString>) -> Result<Self, Error> {
        let py_bytes = bound.clone().unbind();
        let backed_str = PyBackedStr::try_from(bound)?;

        Ok(Self {
            py_bytes,
            backed_str,
        })
    }

    fn as_str(&self) -> &str {
        &self.backed_str
    }

    fn into_py_str(self) -> Py<PyString> {
        self.py_str
    }
}

#[pyfunction]
#[pyo3(signature=(input, input2))]
fn rust_func(
    input: Vec<[Bound<'py, PyString>; 2]>,
    input2: Vec<(Bound<'py, PyString>, i32)>,
) -> PyResult<(Py<PyString>, i32)> {
    let input_unbind = input
        .into_iter()
        .map(|e| {
            let mut e_iter = e.into_iter();
            Ok([
                PyStringWrapper::from_bound(e_iter.next().unwrap())?,
                PyStringWrapper::from_bound(e_iter.next().unwrap())?,
            ])
        })
        .collect::<Result<Vec<_>, Error>>()
        .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;

    let input2_unbind = input2
        .into_iter()
        .map(|e| Ok((PyStringWrapper::from_bound(e.0)?, e.1)))
        .collect::<Result<Vec<_>, Error>>()
        .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;

    let tp = ThreadPoolBuilder::new()
        .num_threads(4 as usize)
        .build()
        .unwrap();

    let res = tp.install(|| {
        input_unbind
            .par_iter()
            .flat_map(|a| {
                input2_unbind
                    .iter()
                    .cloned()
                    .map(move |b| (a, b))
                    .collect::<Vec<_>>()
            })
            .map(|(a, b)| {
                let r = do_something(a);
                (b.0.into_py_str(), r)
            })
            .collect::<Vec<_>>()
    });

    Ok(res)
}