Implement PairedQueue

Didn't have time tonight to get int priorities working in some cases
This commit is contained in:
Michael Bradley 2025-01-09 23:01:59 +13:00
parent 494168597e
commit 7b74ab3687
Signed by: MichaelBradley
SSH key fingerprint: SHA256:cj/YZ5VT+QOKncqSkx+ibKTIn0Obg7OIzwzl9BL8EO8
2 changed files with 94 additions and 7 deletions

View file

@ -10,12 +10,12 @@ pub struct Item<D: Clone, P: PartialOrd + Clone> {
impl<D: Clone, P: PartialOrd + Clone> Item<D, P> {
/// Creates a new instance
fn new(data: D, priority: P) -> Self {
pub fn new(data: D, priority: P) -> Self {
Self { data, priority }
}
/// Retrieve the internal data, it would be nicer to implement this using [`From`] or [`Into`], but I don't see a way to do that using generics
fn data(self) -> D {
pub fn data(self) -> D {
self.data
}
}

View file

@ -1,9 +1,13 @@
// A "pure" priority queue that supports duplicates, but not arbitrary deletions or weight updates
/// 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::prelude::*;
use pyo3::{
exceptions::{PyIndexError, PyStopIteration, PyTypeError},
prelude::*,
types::{PyDict, PyType},
};
#[pyclass]
pub struct PairedQueue {
@ -12,14 +16,97 @@ pub struct PairedQueue {
#[pymethods]
impl PairedQueue {
#[classmethod]
fn __class_getitem__(cls_: Bound<'_, PyType>, _key: Py<PyAny>) -> Bound<'_, PyType> {
cls_
}
#[new]
fn new() -> Self {
Self {
backing: Box::new(BinaryHeap::new()),
#[pyo3(signature = (items=None))]
fn new(items: Option<Py<PyAny>>) -> PyResult<Self> {
if let Some(py_object) = items {
Python::with_gil(|py| Self::new_from_any(py_object.bind(py)))
} 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));
}
}
// TODO: Support u64
impl<'py> PairedQueue {
fn new_from_any(object: &Bound<'py, PyAny>) -> PyResult<Self> {
if let Ok(vec) = object.extract::<Vec<(Py<PyAny>, f64)>>() {
Ok(Self::new_from_vec(vec))
} else {
if object.is_instance_of::<PyDict>() {
if let Ok(dict) = object.downcast::<PyDict>() {
Self::new_from_map(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 new_from_vec(list: Vec<(Py<PyAny>, f64)>) -> Self {
Self {
backing: Box::new(BinaryHeap::from_iter(
list.into_iter()
.map(|(data, priority)| Item::new(data, priority)),
)),
}
}
fn new_from_map(dict: &Bound<'py, PyDict>) -> PyResult<Self> {
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(Self {
backing: Box::new(BinaryHeap::from_iter(items)),
})
} else {
Err(PyErr::new::<PyTypeError, _>("Dict keys were not floats"))
}
}
}