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,
|
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);
|
||||||
|
|
|
@ -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]
|
||||||
|
|
|
@ -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)
|
||||||
|
|
|
@ -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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -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));
|
||||||
|
});
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue