pyo3/tests/test_pyself.rs

124 lines
3.0 KiB
Rust
Raw Normal View History

2021-12-03 00:03:32 +00:00
#![cfg(feature = "macros")]
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};
2022-01-11 07:50:02 +00:00
use pyo3::PyCell;
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
}
2021-02-11 21:37:38 +00:00
fn get_iter(slf: &PyCell<Self>, keys: Py<PyBytes>) -> Iter {
Iter {
reader: slf.into(),
2019-04-03 09:01:28 +00:00
keys,
idx: 0,
2021-02-11 21:37:38 +00:00
}
2019-04-03 09:01:28 +00:00
}
2020-02-15 15:24:38 +00:00
fn get_iter_and_reset(
2022-03-23 07:07:28 +00:00
mut slf: PyRefMut<'_, Self>,
2020-02-15 15:24:38 +00:00
keys: Py<PyBytes>,
2022-03-23 07:07:28 +00:00
py: Python<'_>,
2020-02-15 15:24:38 +00:00
) -> 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,
}
2022-01-11 07:50:02 +00:00
#[pymethods]
impl Iter {
#[allow(clippy::self_named_constructors)]
2022-03-23 07:07:28 +00:00
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}
2019-12-14 14:16:39 +00:00
2022-03-23 07:07:28 +00:00
fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<Option<PyObject>> {
let bytes = slf.keys.as_ref(slf.py()).as_bytes();
2019-12-14 14:16:39 +00:00
match bytes.get(slf.idx) {
Some(&b) => {
2020-02-15 15:24:38 +00:00
slf.idx += 1;
let py = slf.py();
2020-02-15 15:24:38 +00:00
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() {
2022-07-19 17:34:23 +00:00
Python::with_gil(|py| {
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() {
2022-07-19 17:34:23 +00:00
Python::with_gil(|py| {
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()");
});
2019-08-11 06:27:03 +00:00
}
2019-04-03 09:01:28 +00:00
#[test]
fn test_nested_iter_reset() {
2022-07-19 17:34:23 +00:00
Python::with_gil(|py| {
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());
});
2019-04-03 09:01:28 +00:00
}