Add support for set operations with UnionSet:

* difference, intersection, symmetric_difference, union

Author: davidblewett

rust/rust_fst.h

@@ -65,6 +65,7 @@ SetLevStream* fst_set_levsearch(Set*, Levenshtein*);
 SetRegexStream* fst_set_regexsearch(Set*, Regex*);
 SetOpBuilder* fst_set_make_opbuilder(Set*);
 SetOpBuilder* fst_set_make_opbuilder_streambuilder(SetStreamBuilder*);
+SetOpBuilder* fst_set_make_opbuilder_union(SetUnion*);
 void fst_set_free(Set*);
 
 char* fst_set_stream_next(SetStream*);
@@ -78,6 +79,7 @@ void fst_set_regexstream_free(SetRegexStream*);
 
 void fst_set_opbuilder_push(SetOpBuilder*, Set*);
 void fst_set_opbuilder_push_streambuilder(SetOpBuilder*, SetStreamBuilder*);
+void fst_set_opbuilder_push_union(SetOpBuilder*, SetUnion*);
 void fst_set_opbuilder_free(SetOpBuilder*);
 SetUnion* fst_set_opbuilder_union(SetOpBuilder*);
 SetIntersection* fst_set_opbuilder_intersection(SetOpBuilder*);

rust/src/set.rs

@@ -154,6 +154,14 @@ pub extern "C" fn fst_set_make_opbuilder_streambuilder(ptr: *mut set::StreamBuil
     to_raw_ptr(ob)
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_make_opbuilder_union(ptr: *mut set::Union<'static>) -> *mut set::OpBuilder<'static> {
+    let union = val_from_ptr!(ptr);
+    let mut ob = set::OpBuilder::new();
+    ob.push(union.into_stream());
+    to_raw_ptr(ob)
+}
+
 #[no_mangle]
 pub extern "C" fn fst_set_opbuilder_push(ptr: *mut set::OpBuilder, set_ptr: *mut Set) {
     let set = ref_from_ptr!(set_ptr);
@@ -168,6 +176,13 @@ pub extern "C" fn fst_set_opbuilder_push_streambuilder(ptr: *mut set::OpBuilder<
     ob.push(sb.into_stream());
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_opbuilder_push_union(ptr: *mut set::OpBuilder<'static>, union_ptr: *mut set::Union<'static>) {
+    let union = val_from_ptr!(union_ptr);
+    let ob = mutref_from_ptr!(ptr);
+    ob.push(union.into_stream());
+}
+
 #[no_mangle]
 pub extern "C" fn fst_set_opbuilder_union(ptr: *mut set::OpBuilder)
                                           -> *mut set::Union {

rust_fst/set.py

@@ -59,17 +59,20 @@ def get_set(self):
 class OpBuilderInputType(Enum):
     SET = 1
     STREAM_BUILDER = 2
+    UNION = 3
 
 
 class OpBuilder(object):
 
     _BUILDERS = {
         OpBuilderInputType.SET: lib.fst_set_make_opbuilder,
         OpBuilderInputType.STREAM_BUILDER: lib.fst_set_make_opbuilder_streambuilder,
+        OpBuilderInputType.UNION: lib.fst_set_make_opbuilder_union,
     }
     _PUSHERS  =  {
         OpBuilderInputType.SET: lib.fst_set_opbuilder_push,
         OpBuilderInputType.STREAM_BUILDER: lib.fst_set_opbuilder_push_streambuilder,
+        OpBuilderInputType.UNION: lib.fst_set_opbuilder_push_union,
     }
 
     @classmethod
@@ -452,3 +455,66 @@ def __iter__(self):
         for fst in self.sets[1:]:
             opbuilder.push(fst._ptr)
         return opbuilder.union()
+
+    def _make_opbuilder(self, *others):
+        others = list(others)
+        if len(self.sets) <= 1:
+            raise ValueError(
+                "Must have more than one set to operate on.")
+        if not others:
+            raise ValueError(
+                "Must have at least one set to compare against.")
+        our_opbuilder = OpBuilder(self.sets[0]._ptr,
+                                  input_type=OpBuilderInputType.SET)
+        for fst in self.sets[1:]:
+            our_opbuilder.push(fst._ptr)
+        our_stream = lib.fst_set_opbuilder_union(our_opbuilder._ptr)
+
+        their_opbuilder = OpBuilder(others.pop()._ptr,
+                                    input_type=OpBuilderInputType.SET)
+        for fst in others:
+            their_opbuilder.push(fst._ptr)
+        their_stream = lib.fst_set_opbuilder_union(their_opbuilder._ptr)
+
+        opbuilder = OpBuilder(our_stream, input_type=OpBuilderInputType.UNION)
+        opbuilder.push(their_stream)
+        return opbuilder
+
+    def difference(self, *others):
+        """ Get an iterator over the keys in the difference of this set and
+            others.
+
+        :param others:  List of :py:class:`Set` objects
+        :returns:       Iterator over all keys that exists in this set, but in
+                        none of the other sets, in lexicographical order
+        """
+        return self._make_opbuilder(*others).difference()
+
+    def intersection(self, *others):
+        """ Get an iterator over the keys in the intersection of this set and
+            others.
+
+        :param others:  List of :py:class:`Set` objects
+        :returns:       Iterator over all keys that exists in all of the passed
+                        sets in lexicographical order
+        """
+        return self._make_opbuilder(*others).intersection()
+
+    def symmetric_difference(self, *others):
+        """ Get an iterator over the keys in the symmetric difference of this
+            set and others.
+
+        :param others:  List of :py:class:`Set` objects
+        :returns:       Iterator over all keys that exists in only one of the
+                        sets in lexicographical order
+        """
+        return self._make_opbuilder(*others).symmetric_difference()
+
+    def union(self, *others):
+        """ Get an iterator over the keys in the union of this set and others.
+
+        :param others:  List of :py:class:`Set` objects
+        :returns:       Iterator over all keys in all sets in lexicographical
+                        order
+        """
+        return self._make_opbuilder(*others).union()

tests/test_set.py

@@ -166,6 +166,20 @@ def test_unionset_contains(fst_unionset):
         assert key in fst_unionset
 
 
+def test_unionset_difference():
+    a = Set.from_iter(["bar", "foo"])
+    b = Set.from_iter(["baz", "foo"])
+    c = Set.from_iter(["bonk", "foo"])
+    assert list(UnionSet(a, b).difference(c)) == ["bar", "baz"]
+
+
+def test_unionset_intersection():
+    a = Set.from_iter(["bar", "foo"])
+    b = Set.from_iter(["baz", "foo"])
+    c = Set.from_iter(["bonk", "foo"])
+    assert list(UnionSet(a, b).intersection(c)) == ["foo"]
+
+
 def test_unionset_iter(fst_unionset):
     stored_keys = list(fst_unionset)
     assert stored_keys == sorted(set(TEST_KEYS+TEST_KEYS2))
@@ -179,3 +193,17 @@ def test_unionset_range(fst_unionset):
         fst_unionset['c':'a']
     with pytest.raises(ValueError):
         fst_unionset['c']
+
+
+def test_unionset_symmetric_difference():
+    a = Set.from_iter(["bar", "foo"])
+    b = Set.from_iter(["baz", "foo"])
+    c = Set.from_iter(["bonk", "foo"])
+    assert list(UnionSet(a, b).symmetric_difference(c)) == ["bar", "baz", "bonk"]
+
+
+def test_unionset_union():
+    a = Set.from_iter(["bar", "foo"])
+    b = Set.from_iter(["baz", "foo"])
+    c = Set.from_iter(["bonk", "foo"])
+    assert list(UnionSet(a, b).union(c)) == ["bar", "baz", "bonk", "foo"]