110 lines
3.3 KiB
Rust
110 lines
3.3 KiB
Rust
/// A "paired" priority queue that links some data to a priority and supports duplicates, but not arbitrary deletions or weight updates
|
|
use crate::backing::{
|
|
item::Item,
|
|
pure::{BinaryHeap, PureBacking},
|
|
};
|
|
use pyo3::{
|
|
exceptions::{PyIndexError, PyStopIteration, PyTypeError},
|
|
prelude::*,
|
|
types::{PyDict, PyType},
|
|
};
|
|
|
|
#[pyclass]
|
|
pub struct PairedQueue {
|
|
backing: Box<dyn PureBacking<Item<Py<PyAny>, f64>>>,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl PairedQueue {
|
|
#[classmethod]
|
|
fn __class_getitem__(cls_: Bound<'_, PyType>, _key: Py<PyAny>) -> Bound<'_, PyType> {
|
|
cls_
|
|
}
|
|
|
|
#[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))).and_then(|vec| {
|
|
Ok(Self {
|
|
backing: Box::new(BinaryHeap::from_iter(vec)),
|
|
})
|
|
})
|
|
} else {
|
|
Ok(Self {
|
|
backing: Box::new(BinaryHeap::new()),
|
|
})
|
|
}
|
|
}
|
|
|
|
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, _>(()))
|
|
}
|
|
}
|
|
|
|
fn __setitem__(mut self_: PyRefMut<'_, Self>, key: Py<PyAny>, value: f64) {
|
|
self_.backing.add(Item::new(key, value));
|
|
}
|
|
}
|
|
|
|
impl<'py> PairedQueue {
|
|
fn from_any(object: &Bound<'py, PyAny>) -> PyResult<Vec<Item<Py<PyAny>, 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",
|
|
))
|
|
}
|
|
}
|
|
}
|
|
|
|
fn from_vec(list: Vec<(Py<PyAny>, f64)>) -> Vec<Item<Py<PyAny>, f64>> {
|
|
list.into_iter()
|
|
.map(|(data, priority)| Item::new(data, priority))
|
|
.collect()
|
|
}
|
|
|
|
fn from_dict(dict: &Bound<'py, PyDict>) -> PyResult<Vec<Item<Py<PyAny>, f64>>> {
|
|
if let Ok(items) = dict
|
|
.into_iter()
|
|
.map(|(data, priority)| match priority.extract::<f64>() {
|
|
Ok(value) => Ok(Item::new(data.unbind(), value)),
|
|
Err(err) => Err(err),
|
|
})
|
|
.collect::<Result<Vec<_>, _>>()
|
|
{
|
|
Ok(items)
|
|
} else {
|
|
Err(PyErr::new::<PyTypeError, _>("Dict keys were not floats"))
|
|
}
|
|
}
|
|
}
|