pyo3/tests/test_pyself.rs

123 lines
3.1 KiB
Rust

//! Test slf: PyRef/PyMutRef<Self>(especially, slf.into::<Py>) works
use pyo3;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyString};
use pyo3::{AsPyRef, PyCell, PyIterProtocol};
use std::collections::HashMap;
mod common;
/// Assumes it's a file reader or so.
/// Inspired by https://github.com/jothan/cordoba, thanks.
#[pyclass]
#[derive(Clone, Debug)]
struct Reader {
inner: HashMap<u8, String>,
}
#[pymethods]
impl Reader {
fn clone_ref(slf: &PyCell<Self>) -> &PyCell<Self> {
slf
}
fn clone_ref_with_py<'py>(slf: &'py PyCell<Self>, _py: Python<'py>) -> &'py PyCell<Self> {
slf
}
fn get_iter(slf: &PyCell<Self>, keys: Py<PyBytes>) -> PyResult<Iter> {
Ok(Iter {
reader: slf.into(),
keys,
idx: 0,
})
}
fn get_iter_and_reset(
mut slf: PyRefMut<Self>,
keys: Py<PyBytes>,
py: Python,
) -> PyResult<Iter> {
let reader = Py::new(py, slf.clone())?;
slf.inner.clear();
Ok(Iter {
reader,
keys,
idx: 0,
})
}
}
#[pyclass]
#[derive(Debug)]
struct Iter {
reader: Py<Reader>,
keys: Py<PyBytes>,
idx: usize,
}
#[pyproto]
impl PyIterProtocol for Iter {
fn __iter__(slf: PyRefMut<Self>) -> PyResult<PyObject> {
let py = unsafe { Python::assume_gil_acquired() };
Ok(slf.into_py(py))
}
fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<PyObject>> {
let py = unsafe { Python::assume_gil_acquired() };
let bytes = slf.keys.as_ref(py).as_bytes();
match bytes.get(slf.idx) {
Some(&b) => {
slf.idx += 1;
let reader = slf.reader.as_ref(py);
let reader_ref = reader.try_borrow()?;
let res = reader_ref
.inner
.get(&b)
.map(|s| PyString::new(py, s).into());
Ok(res)
}
None => Ok(None),
}
}
}
fn reader() -> Reader {
let reader = [(1, "a"), (2, "b"), (3, "c"), (4, "d"), (5, "e")];
Reader {
inner: reader.iter().map(|(k, v)| (*k, (*v).to_string())).collect(),
}
}
#[test]
fn test_nested_iter() {
let gil = Python::acquire_gil();
let py = gil.python();
let reader: PyObject = reader().into_py(py);
py_assert!(
py,
reader,
"list(reader.get_iter(bytes([3, 5, 2]))) == ['c', 'e', 'b']"
);
}
#[test]
fn test_clone_ref() {
let gil = Python::acquire_gil();
let py = gil.python();
let reader: PyObject = reader().into_py(py);
py_assert!(py, reader, "reader == reader.clone_ref()");
py_assert!(py, reader, "reader == reader.clone_ref_with_py()");
}
#[test]
fn test_nested_iter_reset() {
let gil = Python::acquire_gil();
let py = gil.python();
let reader = PyCell::new(py, reader()).unwrap();
py_assert!(
py,
reader,
"list(reader.get_iter_and_reset(bytes([3, 5, 2]))) == ['c', 'e', 'b']"
);
let reader_ref = reader.borrow();
assert!(reader_ref.inner.is_empty());
}