pyo3/benches/bench_any.rs
Adam Reichold bd7aed4b12 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.
2023-07-09 10:16:20 +02:00

95 lines
2.5 KiB
Rust

use criterion::{criterion_group, criterion_main, Bencher, Criterion};
use pyo3::{
types::{
PyBool, PyByteArray, PyBytes, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PyMapping,
PySequence, PySet, PyString, PyTuple,
},
PyAny, PyResult, Python,
};
#[derive(PartialEq, Eq, Debug)]
enum ObjectType {
None,
Bool,
ByteArray,
Bytes,
Dict,
Float,
FrozenSet,
Int,
List,
Set,
Str,
Tuple,
Sequence,
Mapping,
Unknown,
}
fn find_object_type(obj: &PyAny) -> ObjectType {
if obj.is_none() {
ObjectType::None
} else if obj.is_instance_of::<PyBool>() {
ObjectType::Bool
} else if obj.is_instance_of::<PyByteArray>() {
ObjectType::ByteArray
} else if obj.is_instance_of::<PyBytes>() {
ObjectType::Bytes
} else if obj.is_instance_of::<PyDict>() {
ObjectType::Dict
} else if obj.is_instance_of::<PyFloat>() {
ObjectType::Float
} else if obj.is_instance_of::<PyFrozenSet>() {
ObjectType::FrozenSet
} else if obj.is_instance_of::<PyInt>() {
ObjectType::Int
} else if obj.is_instance_of::<PyList>() {
ObjectType::List
} else if obj.is_instance_of::<PySet>() {
ObjectType::Set
} else if obj.is_instance_of::<PyString>() {
ObjectType::Str
} else if obj.is_instance_of::<PyTuple>() {
ObjectType::Tuple
} else if obj.downcast::<PySequence>().is_ok() {
ObjectType::Sequence
} else if obj.downcast::<PyMapping>().is_ok() {
ObjectType::Mapping
} else {
ObjectType::Unknown
}
}
fn bench_identify_object_type(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
let obj = py.eval("object()", None, None).unwrap();
b.iter(|| find_object_type(obj));
assert_eq!(find_object_type(obj), ObjectType::Unknown);
});
}
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);
criterion_main!(benches);