pyo3/tests/test_pyself.rs

122 lines
3.1 KiB
Rust
Raw Normal View History

2019-04-03 09:01:28 +00:00
//! Test slf: PyRef/PyMutRef<Self>(especially, slf.into::<Py>) works
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyString};
2020-02-03 13:25:16 +00:00
use pyo3::{AsPyRef, PyCell, PyIterProtocol};
use std::collections::HashMap;
mod common;
/// Assumes it's a file reader or so.
2019-04-03 09:01:28 +00:00
/// Inspired by https://github.com/jothan/cordoba, thanks.
#[pyclass]
2019-12-14 14:16:39 +00:00
#[derive(Clone, Debug)]
struct Reader {
inner: HashMap<u8, String>,
}
#[pymethods]
impl Reader {
2020-02-03 13:25:16 +00:00
fn clone_ref(slf: &PyCell<Self>) -> &PyCell<Self> {
2019-08-11 06:27:03 +00:00
slf
}
2020-02-03 13:25:16 +00:00
fn clone_ref_with_py<'py>(slf: &'py PyCell<Self>, _py: Python<'py>) -> &'py PyCell<Self> {
2019-09-01 15:31:22 +00:00
slf
}
2020-02-03 13:25:16 +00:00
fn get_iter(slf: &PyCell<Self>, keys: Py<PyBytes>) -> PyResult<Iter> {
Ok(Iter {
reader: slf.into(),
2019-04-03 09:01:28 +00:00
keys,
idx: 0,
})
}
2020-02-15 15:24:38 +00:00
fn get_iter_and_reset(
mut slf: PyRefMut<Self>,
keys: Py<PyBytes>,
py: Python,
) -> PyResult<Iter> {
2019-04-03 09:01:28 +00:00
let reader = Py::new(py, slf.clone())?;
slf.inner.clear();
Ok(Iter {
reader,
keys,
idx: 0,
})
}
}
#[pyclass]
2020-02-15 15:24:38 +00:00
#[derive(Debug)]
struct Iter {
reader: Py<Reader>,
keys: Py<PyBytes>,
idx: usize,
}
#[pyproto]
impl PyIterProtocol for Iter {
fn __iter__(slf: PyRef<Self>) -> PyResult<PyObject> {
let py = unsafe { Python::assume_gil_acquired() };
2020-02-15 15:24:38 +00:00
Ok(slf.into_py(py))
}
2019-12-14 14:16:39 +00:00
2020-02-15 15:24:38 +00:00
fn __next__(mut slf: PyRefMut<Self>) -> PyResult<Option<PyObject>> {
let py = unsafe { Python::assume_gil_acquired() };
2019-12-14 14:16:39 +00:00
let bytes = slf.keys.as_ref(py).as_bytes();
match bytes.get(slf.idx) {
Some(&b) => {
2020-02-15 15:24:38 +00:00
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),
}
}
}
2019-08-11 06:27:03 +00:00
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(),
2019-08-11 06:27:03 +00:00
}
}
#[test]
2020-02-15 15:24:38 +00:00
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']"
);
}
2019-04-03 09:01:28 +00:00
2019-08-11 06:27:03 +00:00
#[test]
fn test_clone_ref() {
let gil = Python::acquire_gil();
let py = gil.python();
let reader: PyObject = reader().into_py(py);
2019-08-11 06:27:03 +00:00
py_assert!(py, reader, "reader == reader.clone_ref()");
2019-09-01 15:31:22 +00:00
py_assert!(py, reader, "reader == reader.clone_ref_with_py()");
2019-08-11 06:27:03 +00:00
}
2019-04-03 09:01:28 +00:00
#[test]
fn test_nested_iter_reset() {
let gil = Python::acquire_gil();
let py = gil.python();
2020-02-15 15:24:38 +00:00
let reader = PyCell::new(py, reader()).unwrap();
2019-04-03 09:01:28 +00:00
py_assert!(
py,
2019-08-11 06:27:03 +00:00
reader,
"list(reader.get_iter_and_reset(bytes([3, 5, 2]))) == ['c', 'e', 'b']"
2019-04-03 09:01:28 +00:00
);
2020-02-15 15:24:38 +00:00
let reader_ref = reader.borrow();
assert!(reader_ref.inner.is_empty());
2019-04-03 09:01:28 +00:00
}