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:
parent
26018560f2
commit
bd7aed4b12
|
@ -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);
|
||||
|
|
|
@ -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]
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -122,7 +122,7 @@ impl<T> Default for PyClassImplCollector<T> {
|
|||
|
||||
impl<T> Clone for PyClassImplCollector<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self::new()
|
||||
*self
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
@ -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));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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)
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue