//! Test slf: PyRef/PyMutRef(especially, slf.into::) 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, } #[pymethods] impl Reader { fn clone_ref(slf: &PyCell) -> &PyCell { slf } fn clone_ref_with_py<'py>(slf: &'py PyCell, _py: Python<'py>) -> &'py PyCell { slf } fn get_iter(slf: &PyCell, keys: Py) -> PyResult { Ok(Iter { reader: slf.into(), keys, idx: 0, }) } fn get_iter_and_reset( mut slf: PyRefMut, keys: Py, py: Python, ) -> PyResult { let reader = Py::new(py, slf.clone())?; slf.inner.clear(); Ok(Iter { reader, keys, idx: 0, }) } } #[pyclass] #[derive(Debug)] struct Iter { reader: Py, keys: Py, idx: usize, } #[pyproto] impl PyIterProtocol for Iter { fn __iter__(slf: PyRefMut) -> PyResult { let py = unsafe { Python::assume_gil_acquired() }; Ok(slf.into_py(py)) } fn __next__(mut slf: PyRefMut) -> PyResult> { 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()); }