Add rpds's Stack.

Author: Julian

src/lib.rs

@@ -5,6 +5,7 @@ use pyo3::{exceptions::PyKeyError, types::PyMapping, types::PyTupleMethods};
 use pyo3::{prelude::*, BoundObject, PyTypeInfo};
 use rpds::{
     HashTrieMap, HashTrieMapSync, HashTrieSet, HashTrieSetSync, List, ListSync, Queue, QueueSync,
+    Stack, StackSync,
 };
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};
@@ -1238,6 +1239,155 @@ impl ListIterator {
     }
 }
 
+#[repr(transparent)]
+#[pyclass(name = "Stack", module = "rpds", frozen, sequence)]
+struct StackPy {
+    inner: StackSync<Py<PyAny>>,
+}
+
+impl From<StackSync<Py<PyAny>>> for StackPy {
+    fn from(elements: StackSync<Py<PyAny>>) -> Self {
+        StackPy { inner: elements }
+    }
+}
+
+#[pymethods]
+impl StackPy {
+    #[new]
+    #[pyo3(signature = (*args))]
+    fn init(args: &Bound<'_, PyTuple>) -> PyResult<Self> {
+        let mut inner = Stack::new_sync();
+        if args.len() == 1 {
+            for each in args.get_item(0)?.try_iter()? {
+                inner.push_mut(each?.extract()?);
+            }
+        } else {
+            for each in args {
+                inner.push_mut(each.extract()?);
+            }
+        }
+        Ok(StackPy { inner })
+    }
+
+    fn __hash__(&self, py: Python<'_>) -> PyResult<u64> {
+        let mut hasher = DefaultHasher::new();
+
+        self.inner
+            .iter()
+            .enumerate()
+            .try_for_each(|(index, each)| {
+                each.bind(py)
+                    .hash()
+                    .map_err(|_| {
+                        PyTypeError::new_err(format!(
+                            "Unhashable type at {} element in Stack: {}",
+                            index,
+                            each.bind(py)
+                                .repr()
+                                .and_then(|r| r.extract())
+                                .unwrap_or("<repr> error".to_string())
+                        ))
+                    })
+                    .map(|x| hasher.write_isize(x))
+            })?;
+
+        Ok(hasher.finish())
+    }
+
+    fn __iter__(slf: PyRef<'_, Self>) -> StackIterator {
+        StackIterator {
+            inner: slf.inner.clone(),
+        }
+    }
+
+    fn __len__(&self) -> usize {
+        self.inner.size()
+    }
+
+    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyResult<Py<PyAny>> {
+        match op {
+            CompareOp::Eq => (self.inner.size() == other.inner.size()
+                && self
+                    .inner
+                    .iter()
+                    .zip(other.inner.iter())
+                    .map(|(e1, e2)| e1.bind(py).eq(e2))
+                    .all(|r| r.unwrap_or(false)))
+            .into_pyobject(py)
+            .map_err(Into::into)
+            .map(BoundObject::into_any)
+            .map(BoundObject::unbind),
+            CompareOp::Ne => (self.inner.size() != other.inner.size()
+                || self
+                    .inner
+                    .iter()
+                    .zip(other.inner.iter())
+                    .map(|(e1, e2)| e1.bind(py).ne(e2))
+                    .any(|r| r.unwrap_or(true)))
+            .into_pyobject(py)
+            .map_err(Into::into)
+            .map(BoundObject::into_any)
+            .map(BoundObject::unbind),
+            _ => Ok(py.NotImplemented()),
+        }
+    }
+
+    fn __repr__(&self, py: Python) -> PyResult<String> {
+        let contents = self.inner.into_iter().map(|k| {
+            Ok(k.into_pyobject(py)?
+                .call_method0("__repr__")
+                .and_then(|r| r.extract())
+                .unwrap_or("<repr failed>".to_owned()))
+        });
+        let mut contents = contents.collect::<Result<Vec<_>, PyErr>>()?;
+        contents.reverse();
+        Ok(format!("Stack([{}])", contents.join(", ")))
+    }
+
+    fn peek(&self, py: Python) -> PyResult<Py<PyAny>> {
+        if let Some(peeked) = self.inner.peek() {
+            Ok(peeked.clone_ref(py))
+        } else {
+            Err(PyIndexError::new_err("peeked an empty stack"))
+        }
+    }
+
+    fn pop(&self) -> PyResult<StackPy> {
+        if let Some(popped) = self.inner.pop() {
+            Ok(StackPy { inner: popped })
+        } else {
+            Err(PyIndexError::new_err("popped an empty stack"))
+        }
+    }
+
+    fn push(&self, other: Py<PyAny>) -> StackPy {
+        StackPy {
+            inner: self.inner.push(other),
+        }
+    }
+}
+
+#[pyclass(module = "rpds")]
+struct StackIterator {
+    inner: StackSync<Py<PyAny>>,
+}
+
+#[pymethods]
+impl StackIterator {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<Py<PyAny>> {
+        let first_op = slf.inner.peek()?;
+        let first = first_op.clone_ref(slf.py());
+
+        slf.inner = slf.inner.pop()?;
+
+        Some(first)
+    }
+}
+
 #[pyclass(module = "rpds")]
 struct QueueIterator {
     inner: QueueSync<Py<PyAny>>,
@@ -1401,6 +1551,7 @@ fn rpds_py(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<HashTrieMapPy>()?;
     m.add_class::<HashTrieSetPy>()?;
     m.add_class::<ListPy>()?;
+    m.add_class::<StackPy>()?;
     m.add_class::<QueuePy>()?;
 
     PyMapping::register::<HashTrieMapPy>(py)?;

tests/test_stack.py

@@ -0,0 +1,151 @@
+"""
+Modified from the pyrsistent test suite.
+
+Pre-modification, these were MIT licensed, and are copyright:
+
+    Copyright (c) 2022 Tobias Gustafsson
+
+    Permission is hereby granted, free of charge, to any person
+    obtaining a copy of this software and associated documentation
+    files (the "Software"), to deal in the Software without
+    restriction, including without limitation the rights to use,
+    copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the
+    Software is furnished to do so, subject to the following
+    conditions:
+
+    The above copyright notice and this permission notice shall be
+    included in all copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+    OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+    NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+    HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+    WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+    OTHER DEALINGS IN THE SOFTWARE.
+"""
+
+import pytest
+
+from rpds import Stack
+
+
+def test_literalish_works():
+    assert Stack(1, 2, 3) == Stack([1, 2, 3])
+
+
+def test_pop_and_peek():
+    ps = Stack([1, 2])
+    assert ps.peek() == 2
+    assert ps.pop().peek() == 1
+    assert ps.pop().pop() == Stack()
+
+
+def test_instantiate_large_stack():
+    assert Stack(range(1000)).peek() == 999
+
+
+def test_iteration():
+    assert list(Stack()) == []
+    assert list(Stack([1, 2, 3]))[::-1] == [1, 2, 3]
+
+
+def test_push():
+    assert Stack([1, 2, 3]).push(4) == Stack([1, 2, 3, 4])
+
+
+def test_push_empty_stack():
+    assert Stack().push(0) == Stack([0])
+
+
+def test_truthiness():
+    assert Stack([1])
+    assert not Stack()
+
+
+def test_len():
+    assert len(Stack([1, 2, 3])) == 3
+    assert len(Stack()) == 0
+
+
+def test_peek_illegal_on_empty_stack():
+    with pytest.raises(IndexError):
+        Stack().peek()
+
+
+def test_pop_illegal_on_empty_stack():
+    with pytest.raises(IndexError):
+        Stack().pop()
+
+
+def test_inequality():
+    assert Stack([1, 2]) != Stack([1, 3])
+    assert Stack([1, 2]) != Stack([1, 2, 3])
+    assert Stack() != Stack([1, 2, 3])
+
+
+def test_repr():
+    assert str(Stack()) == "Stack([])"
+    assert str(Stack([1, 2, 3])) in "Stack([1, 2, 3])"
+
+
+def test_hashing():
+    o = object()
+
+    assert hash(Stack([o, o])) == hash(Stack([o, o]))
+    assert hash(Stack([o])) == hash(Stack([o]))
+    assert hash(Stack()) == hash(Stack([]))
+    assert not (hash(Stack([1, 2])) == hash(Stack([1, 3])))
+    assert not (hash(Stack([1, 2])) == hash(Stack([2, 1])))
+    assert not (hash(Stack([o])) == hash(Stack([o, o])))
+    assert not (hash(Stack([])) == hash(Stack([o])))
+
+    assert hash(Stack([1, 2])) != hash(Stack([1, 3]))
+    assert hash(Stack([1, 2])) != hash(Stack([2, 1]))
+    assert hash(Stack([o])) != hash(Stack([o, o]))
+    assert hash(Stack([])) != hash(Stack([o]))
+    assert not (hash(Stack([o, o])) != hash(Stack([o, o])))
+    assert not (hash(Stack([o])) != hash(Stack([o])))
+    assert not (hash(Stack([])) != hash(Stack([])))
+
+
+def test_sequence():
+    m = Stack("asdf")
+    assert m == Stack(["a", "s", "d", "f"])
+
+
+# Non-pyrsistent-test-suite tests
+
+
+def test_more_eq():
+    o = object()
+
+    assert Stack([o, o]) == Stack([o, o])
+    assert Stack([o]) == Stack([o])
+    assert Stack() == Stack([])
+    assert not (Stack([1, 2]) == Stack([1, 3]))
+    assert not (Stack([o]) == Stack([o, o]))
+    assert not (Stack([]) == Stack([o]))
+
+    assert Stack([1, 2]) != Stack([1, 3])
+    assert Stack([o]) != Stack([o, o])
+    assert Stack([]) != Stack([o])
+    assert not (Stack([o, o]) != Stack([o, o]))
+    assert not (Stack([o]) != Stack([o]))
+    assert not (Stack() != Stack([]))
+
+
+def test_rpds_doc():
+    """
+    From the rpds docs.
+    """
+    stack = Stack().push("stack")
+    assert stack.peek() == "stack"
+
+    a_stack = stack.push("a")
+    assert a_stack.peek() == "a"
+
+    stack_popped = a_stack.pop()
+    assert stack_popped == stack