126 lines
4 KiB
Rust
126 lines
4 KiB
Rust
use pyo3::{
|
|
exceptions::{PyIndexError, PyKeyError, PyStopIteration, PyTypeError},
|
|
prelude::*,
|
|
types::{PyDict, PyType},
|
|
};
|
|
|
|
use crate::backing::{
|
|
containers::PyItem,
|
|
indexed::{IndexedBacking, IndexedBinaryHeap},
|
|
};
|
|
|
|
#[pyclass]
|
|
pub struct IndexedQueue {
|
|
backing: Box<dyn IndexedBacking<PyItem, f64>>,
|
|
}
|
|
|
|
// TODO: Too much of this is just too similar to the other queues
|
|
// Normally I'd solve this with inheritance, but I'm not so familiar with rust trait default implementations yet
|
|
// In my limited understanding they don't seem to play nice with the Pyo3 bindings
|
|
#[pymethods]
|
|
impl IndexedQueue {
|
|
#[new]
|
|
#[pyo3(signature = (items=None))]
|
|
fn new(items: Option<Py<PyAny>>) -> PyResult<Self> {
|
|
if let Some(py_object) = items {
|
|
Python::with_gil(|py| Self::from_any(py_object.bind(py))).map(|vec| Self {
|
|
backing: Box::new(IndexedBinaryHeap::from_iter(vec)),
|
|
})
|
|
} else {
|
|
Ok(Self {
|
|
backing: Box::new(IndexedBinaryHeap::new()),
|
|
})
|
|
}
|
|
}
|
|
|
|
fn __setitem__(mut self_: PyRefMut<'_, Self>, key: Py<PyAny>, value: f64) {
|
|
self_.backing.set(PyItem::new(key), value)
|
|
}
|
|
|
|
fn __delitem__(mut self_: PyRefMut<'_, Self>, key: Py<PyAny>) -> PyResult<()> {
|
|
match self_.backing.remove(PyItem::new(key)) {
|
|
Some(_) => Ok(()),
|
|
None => Err(PyErr::new::<PyKeyError, _>(
|
|
"Key was not contained in queue",
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn __contains__(self_: PyRef<'_, Self>, key: Py<PyAny>) -> bool {
|
|
self_.backing.contains(&PyItem::new(key))
|
|
}
|
|
|
|
/// Enables generic typing
|
|
#[classmethod]
|
|
fn __class_getitem__(cls_: Bound<'_, PyType>, _key: Py<PyAny>) -> Bound<'_, PyType> {
|
|
cls_
|
|
}
|
|
|
|
fn __len__(self_: PyRef<'_, Self>) -> usize {
|
|
self_.backing.len()
|
|
}
|
|
|
|
fn __iter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
|
self_
|
|
}
|
|
|
|
fn __next__(mut self_: PyRefMut<'_, Self>) -> PyResult<Py<PyAny>> {
|
|
if let Some(item) = self_.backing.pop() {
|
|
Ok(item.data())
|
|
} else {
|
|
Err(PyErr::new::<PyStopIteration, _>(()))
|
|
}
|
|
}
|
|
|
|
fn pop(mut self_: PyRefMut<'_, Self>) -> PyResult<Py<PyAny>> {
|
|
if let Some(item) = self_.backing.pop() {
|
|
Ok(item.data())
|
|
} else {
|
|
Err(PyErr::new::<PyIndexError, _>(()))
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<'py> IndexedQueue {
|
|
/// Tries to a Python object into a vector suitable for ingestion into the backing
|
|
fn from_any(object: &Bound<'py, PyAny>) -> PyResult<Vec<(PyItem, f64)>> {
|
|
if let Ok(vec) = object.extract::<Vec<(Py<PyAny>, f64)>>() {
|
|
Ok(Self::from_vec(vec))
|
|
} else if object.is_instance_of::<PyDict>() {
|
|
if let Ok(dict) = object.downcast::<PyDict>() {
|
|
Self::from_dict(dict)
|
|
} else {
|
|
Err(PyErr::new::<PyTypeError, _>(
|
|
"Argument claimed to be a dict but wasn't",
|
|
))
|
|
}
|
|
} else {
|
|
Err(PyErr::new::<PyTypeError, _>(
|
|
"Argument was not a properly-formed dict, list, or tuple",
|
|
))
|
|
}
|
|
}
|
|
|
|
/// Converts a vector of Python objects and priorities into a vector of items
|
|
fn from_vec(list: Vec<(Py<PyAny>, f64)>) -> Vec<(PyItem, f64)> {
|
|
list.into_iter()
|
|
.map(|(data, priority)| (PyItem::new(data), priority))
|
|
.collect()
|
|
}
|
|
|
|
/// Converts a Python dictionary into a vector of items
|
|
fn from_dict(dict: &Bound<'py, PyDict>) -> PyResult<Vec<(PyItem, f64)>> {
|
|
if let Ok(items) = dict
|
|
.into_iter()
|
|
.map(|(data, priority)| match priority.extract::<f64>() {
|
|
Ok(value) => Ok((PyItem::new(data.unbind()), value)),
|
|
Err(err) => Err(err),
|
|
})
|
|
.collect::<Result<Vec<_>, _>>()
|
|
{
|
|
Ok(items)
|
|
} else {
|
|
Err(PyErr::new::<PyTypeError, _>("Dict keys were not numbers"))
|
|
}
|
|
}
|
|
}
|