From 7b74ab3687d165a8e267e5bc56fffd35a2262207 Mon Sep 17 00:00:00 2001
From: Michael Bradley <michaelmbrad@gmail.com>
Date: Thu, 9 Jan 2025 23:01:59 +1300
Subject: [PATCH] Implement PairedQueue

Didn't have time tonight to get int priorities working in some cases
---
 src/backing/item.rs |  4 +-
 src/queue/paired.rs | 97 ++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 94 insertions(+), 7 deletions(-)

diff --git a/src/backing/item.rs b/src/backing/item.rs
index 55ac901..8c926a5 100644
--- a/src/backing/item.rs
+++ b/src/backing/item.rs
@@ -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
     }
 }
diff --git a/src/queue/paired.rs b/src/queue/paired.rs
index ab837cf..c8cda9f 100644
--- a/src/queue/paired.rs
+++ b/src/queue/paired.rs
@@ -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"))
+        }
+    }
 }