Merge pull request #3303 from PyO3/pyiter-size-hint

RFC: Add implementation of Iterator::size_hint for PyIterator
This commit is contained in:
Adam Reichold 2023-07-09 11:45:49 +00:00 committed by GitHub
commit db960e5fb4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 38 additions and 5 deletions

View File

@ -5,7 +5,7 @@ use pyo3::{
PyBool, PyByteArray, PyBytes, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PyMapping,
PySequence, PySet, PyString, PyTuple,
},
PyAny, Python,
PyAny, PyResult, Python,
};
#[derive(PartialEq, Eq, Debug)]
@ -71,8 +71,23 @@ fn bench_identify_object_type(b: &mut Bencher<'_>) {
});
}
fn bench_collect_generic_iterator(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
let collection = py.eval("list(range(1 << 20))", None, None).unwrap();
b.iter(|| {
collection
.iter()
.unwrap()
.collect::<PyResult<Vec<_>>>()
.unwrap()
});
});
}
fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("identify_object_type", bench_identify_object_type);
c.bench_function("collect_generic_iterator", bench_collect_generic_iterator);
}
criterion_group!(benches, criterion_benchmark);

View File

@ -24,12 +24,12 @@ extern "C" {
#[inline]
pub unsafe fn Py_False() -> *mut PyObject {
addr_of_mut!(_Py_FalseStruct) as *mut PyLongObject as *mut PyObject
addr_of_mut!(_Py_FalseStruct) as *mut PyObject
}
#[inline]
pub unsafe fn Py_True() -> *mut PyObject {
addr_of_mut!(_Py_TrueStruct) as *mut PyLongObject as *mut PyObject
addr_of_mut!(_Py_TrueStruct) as *mut PyObject
}
#[inline]

View File

@ -179,7 +179,7 @@ pub fn take_pyo3_options<T: Parse>(attrs: &mut Vec<syn::Attribute>) -> Result<Ve
let mut out = Vec::new();
take_attributes(attrs, |attr| {
if let Some(options) = get_pyo3_options(attr)? {
out.extend(options.into_iter());
out.extend(options);
Ok(true)
} else {
Ok(false)

View File

@ -122,7 +122,7 @@ impl<T> Default for PyClassImplCollector<T> {
impl<T> Clone for PyClassImplCollector<T> {
fn clone(&self) -> Self {
Self::new()
*self
}
}

View File

@ -55,6 +55,12 @@ impl<'p> Iterator for &'p PyIterator {
None => PyErr::take(py).map(Err),
}
}
#[cfg(not(Py_LIMITED_API))]
fn size_hint(&self) -> (usize, Option<usize>) {
let hint = unsafe { ffi::PyObject_LengthHint(self.0.as_ptr(), 0) };
(hint.max(0) as usize, None)
}
}
// PyIter_Check does not exist in the limited API until 3.8
@ -313,4 +319,15 @@ def fibonacci(target):
);
});
}
#[test]
#[cfg(not(Py_LIMITED_API))]
fn length_hint_becomes_size_hint_lower_bound() {
Python::with_gil(|py| {
let list = py.eval("[1, 2, 3]", None, None).unwrap();
let iter = list.iter().unwrap();
let hint = iter.size_hint();
assert_eq!(hint, (3, None));
});
}
}

View File

@ -196,6 +196,7 @@ macro_rules! pyobject_native_type_info(
const MODULE: ::std::option::Option<&'static str> = $module;
#[inline]
#[allow(clippy::redundant_closure_call)]
fn type_object_raw(py: $crate::Python<'_>) -> *mut $crate::ffi::PyTypeObject {
$typeobject(py)
}