use err::{PyDowncastError, PyErr, PyResult};
use ffi;
use instance::PyObjectWithToken;
use objects::PyObjectRef;
use python::{Python, ToPyPointer};
pub struct PyIterator<'p>(&'p PyObjectRef);
impl<'p> PyIterator<'p> {
pub fn from_object<T>(py: Python<'p>, obj: &T) -> Result<PyIterator<'p>, PyDowncastError>
where
T: ToPyPointer,
{
unsafe {
let ptr = ffi::PyObject_GetIter(obj.as_ptr());
if ffi::PyIter_Check(ptr) != 0 {
Ok(PyIterator(py.from_borrowed_ptr(ptr)))
} else {
Err(PyDowncastError)
}
}
}
}
impl<'p> Iterator for PyIterator<'p> {
type Item = PyResult<&'p PyObjectRef>;
fn next(&mut self) -> Option<Self::Item> {
let py = self.0.py();
match unsafe { py.from_owned_ptr_or_opt(ffi::PyIter_Next(self.0.as_ptr())) } {
Some(obj) => Some(Ok(obj)),
None => {
if PyErr::occurred(py) {
Some(Err(PyErr::fetch(py)))
} else {
None
}
}
}
}
}
impl<'p> Drop for PyIterator<'p> {
fn drop(&mut self) {
unsafe { ffi::Py_DECREF(self.0.as_ptr()) }
}
}
#[cfg(test)]
mod tests {
use conversion::{PyTryFrom, ToPyObject};
use instance::AsPyRef;
use objectprotocol::ObjectProtocol;
use objects::{PyList, PyObjectRef};
use python::Python;
use pythonrun::GILPool;
#[test]
fn vec_iter() {
let gil_guard = Python::acquire_gil();
let py = gil_guard.python();
let obj = vec![10, 20].to_object(py);
let inst = <PyObjectRef as PyTryFrom>::try_from(obj.as_ref(py)).unwrap();
let mut it = inst.iter().unwrap();
assert_eq!(10, it.next().unwrap().unwrap().extract().unwrap());
assert_eq!(20, it.next().unwrap().unwrap().extract().unwrap());
assert!(it.next().is_none());
}
#[test]
fn iter_refcnt() {
let obj;
let count;
{
let gil_guard = Python::acquire_gil();
let py = gil_guard.python();
obj = vec![10, 20].to_object(py);
count = obj.get_refcnt();
}
{
let gil_guard = Python::acquire_gil();
let py = gil_guard.python();
let inst = <PyObjectRef as PyTryFrom>::try_from(obj.as_ref(py)).unwrap();
let mut it = inst.iter().unwrap();
assert_eq!(10, it.next().unwrap().unwrap().extract().unwrap());
}
assert_eq!(count, obj.get_refcnt());
}
#[test]
fn iter_item_refcnt() {
let gil_guard = Python::acquire_gil();
let py = gil_guard.python();
let obj;
let none;
let count;
{
let _pool = GILPool::new();
let l = PyList::empty(py);
none = py.None();
l.append(10).unwrap();
l.append(&none).unwrap();
count = none.get_refcnt();
obj = l.to_object(py);
}
{
let _pool = GILPool::new();
let inst = <PyObjectRef as PyTryFrom>::try_from(obj.as_ref(py)).unwrap();
let mut it = inst.iter().unwrap();
assert_eq!(10, it.next().unwrap().unwrap().extract().unwrap());
assert!(it.next().unwrap().unwrap().is_none());
}
assert_eq!(count, none.get_refcnt());
}
}