Add implementation of Iterator::size_hint for PyIterator

When the Python iterator backing `PyIterator` has a `__length_hint__` special
method, we can use this as a lower bound for Rust's `Iterator::size_hint` to
e.g. support pre-allocation of collections.

This is implemented using `PyObject_LengthHint` which is not available in the
stable ABI and hence so is `Iterator::size_hint`. This should be fine since this
is an optimization in any case and the stable ABI is expected to have slightly
worse performance overall.
This commit is contained in:
Adam Reichold 2023-07-08 21:49:18 +02:00
parent 26018560f2
commit bd7aed4b12
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, PyBool, PyByteArray, PyBytes, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PyMapping,
PySequence, PySet, PyString, PyTuple, PySequence, PySet, PyString, PyTuple,
}, },
PyAny, Python, PyAny, PyResult, Python,
}; };
#[derive(PartialEq, Eq, Debug)] #[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) { fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("identify_object_type", bench_identify_object_type); 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); criterion_group!(benches, criterion_benchmark);

View file

@ -24,12 +24,12 @@ extern "C" {
#[inline] #[inline]
pub unsafe fn Py_False() -> *mut PyObject { 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] #[inline]
pub unsafe fn Py_True() -> *mut PyObject { 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] #[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(); let mut out = Vec::new();
take_attributes(attrs, |attr| { take_attributes(attrs, |attr| {
if let Some(options) = get_pyo3_options(attr)? { if let Some(options) = get_pyo3_options(attr)? {
out.extend(options.into_iter()); out.extend(options);
Ok(true) Ok(true)
} else { } else {
Ok(false) Ok(false)

View file

@ -122,7 +122,7 @@ impl<T> Default for PyClassImplCollector<T> {
impl<T> Clone for PyClassImplCollector<T> { impl<T> Clone for PyClassImplCollector<T> {
fn clone(&self) -> Self { 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), 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 // 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; const MODULE: ::std::option::Option<&'static str> = $module;
#[inline] #[inline]
#[allow(clippy::redundant_closure_call)]
fn type_object_raw(py: $crate::Python<'_>) -> *mut $crate::ffi::PyTypeObject { fn type_object_raw(py: $crate::Python<'_>) -> *mut $crate::ffi::PyTypeObject {
$typeobject(py) $typeobject(py)
} }