HashTrieSet union, intersection, difference, symm diff

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rpds-py"
-version = "0.5.3"
+version = "0.6.0"
 edition = "2021"
 
 [lib]

src/lib.rs

@@ -285,6 +285,22 @@ impl HashTrieSetPy {
         }
     }
 
+    fn __and__(&self, other: &Self) -> Self {
+        self.intersection(&other)
+    }
+
+    fn __or__(&self, other: &Self) -> Self {
+        self.union(&other)
+    }
+
+    fn __sub__(&self, other: &Self) -> Self {
+        self.difference(&other)
+    }
+
+    fn __xor__(&self, other: &Self) -> Self {
+        self.symmetric_difference(&other)
+    }
+
     fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<KeyIterator>> {
         let iter = slf
             .inner
@@ -350,6 +366,69 @@ impl HashTrieSetPy {
         }
     }
 
+    fn difference(&self, other: &Self) -> HashTrieSetPy {
+        let mut inner = self.inner.clone();
+        for value in other.inner.iter() {
+            inner.remove_mut(value);
+        }
+        HashTrieSetPy { inner }
+    }
+
+    fn intersection(&self, other: &Self) -> HashTrieSetPy {
+        let mut inner: HashTrieSet<Key> = HashTrieSet::new();
+        let larger: &HashTrieSet<Key>;
+        let iter;
+        if self.inner.size() > other.inner.size() {
+            larger = &self.inner;
+            iter = other.inner.iter();
+        } else {
+            larger = &other.inner;
+            iter = self.inner.iter();
+        }
+        for value in iter {
+            if larger.contains(value) {
+                inner.insert_mut(value.to_owned());
+            }
+        }
+        HashTrieSetPy { inner }
+    }
+
+    fn symmetric_difference(&self, other: &Self) -> HashTrieSetPy {
+        let mut inner: HashTrieSet<Key>;
+        let iter;
+        if self.inner.size() > other.inner.size() {
+            inner = self.inner.clone();
+            iter = other.inner.iter();
+        } else {
+            inner = other.inner.clone();
+            iter = self.inner.iter();
+        }
+        for value in iter {
+            if inner.contains(value) {
+                inner.remove_mut(value);
+            } else {
+                inner.insert_mut(value.to_owned());
+            }
+        }
+        HashTrieSetPy { inner }
+    }
+
+    fn union(&self, other: &Self) -> HashTrieSetPy {
+        let mut inner: HashTrieSet<Key>;
+        let iter;
+        if self.inner.size() > other.inner.size() {
+            inner = self.inner.clone();
+            iter = other.inner.iter();
+        } else {
+            inner = other.inner.clone();
+            iter = self.inner.iter();
+        }
+        for value in iter {
+            inner.insert_mut(value.to_owned());
+        }
+        HashTrieSetPy { inner }
+    }
+
     #[pyo3(signature = (*iterables))]
     fn update(&self, iterables: &PyTuple) -> PyResult<HashTrieSetPy> {
         let mut inner = self.inner.clone();

tests/test_hash_trie_set.py

@@ -95,7 +95,6 @@ def test_contains():
     assert 4 not in s
 
 
-@pytest.mark.xfail(reason="Can't figure out inheriting collections.abc yet")
 def test_supports_set_operations():
     s1 = HashTrieSet([1, 2, 3])
     s2 = HashTrieSet([3, 4, 5])