Add UnionSet class:

* Supports range queries across multiple sets:
  a = Set.from_iter(["bar", "foo"])
  b = Set.from_iter(["baz", "foo"])
  list(UnionSet(a, b)['ba':'bb'])
  ['bar', 'baz']
* Add StreamBuilder, to correspond with the Rust library's abstraction
* Update OpBuilder to support constructing operations against multiple
  underlying types (Set and StreamBuilder for now)

Author: davidblewett

rust/rust_fst.h

@@ -64,6 +64,7 @@ SetStream* fst_set_stream(Set*);
 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*);
 void fst_set_free(Set*);
 
 char* fst_set_stream_next(SetStream*);
@@ -76,6 +77,7 @@ char* fst_set_regexstream_next(SetRegexStream*);
 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_free(SetOpBuilder*);
 SetUnion* fst_set_opbuilder_union(SetOpBuilder*);
 SetIntersection* fst_set_opbuilder_intersection(SetOpBuilder*);
@@ -97,6 +99,8 @@ void fst_set_symmetricdifference_free(SetSymmetricDifference*);
 
 SetStreamBuilder* fst_set_streambuilder_new(Set*);
 SetStreamBuilder* fst_set_streambuilder_add_ge(SetStreamBuilder*, char*);
+SetStreamBuilder* fst_set_streambuilder_add_gt(SetStreamBuilder*, char*);
+SetStreamBuilder* fst_set_streambuilder_add_le(SetStreamBuilder*, char*);
 SetStreamBuilder* fst_set_streambuilder_add_lt(SetStreamBuilder*, char*);
 SetStream* fst_set_streambuilder_finish(SetStreamBuilder*);
 

rust/src/set.rs

@@ -146,13 +146,28 @@ pub extern "C" fn fst_set_make_opbuilder(ptr: *mut Set) -> *mut set::OpBuilder<'
 }
 make_free_fn!(fst_set_opbuilder_free, *mut set::OpBuilder);
 
+#[no_mangle]
+pub extern "C" fn fst_set_make_opbuilder_streambuilder(ptr: *mut set::StreamBuilder<'static>) -> *mut set::OpBuilder<'static> {
+    let sb = val_from_ptr!(ptr);
+    let mut ob = set::OpBuilder::new();
+    ob.push(sb.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);
     let ob = mutref_from_ptr!(ptr);
     ob.push(set);
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_opbuilder_push_streambuilder(ptr: *mut set::OpBuilder<'static>, sb_ptr: *mut set::StreamBuilder<'static>) {
+    let sb = val_from_ptr!(sb_ptr);
+    let ob = mutref_from_ptr!(ptr);
+    ob.push(sb.into_stream());
+}
+
 #[no_mangle]
 pub extern "C" fn fst_set_opbuilder_union(ptr: *mut set::OpBuilder)
                                           -> *mut set::Union {
@@ -205,6 +220,22 @@ pub extern "C" fn fst_set_streambuilder_add_ge(ptr: *mut set::StreamBuilder<'sta
     to_raw_ptr(sb.ge(cstr_to_str(c_bound)))
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_streambuilder_add_gt(ptr: *mut set::StreamBuilder<'static>,
+                                               c_bound: *mut libc::c_char)
+                                               -> *mut set::StreamBuilder<'static> {
+    let sb = val_from_ptr!(ptr);
+    to_raw_ptr(sb.gt(cstr_to_str(c_bound)))
+}
+
+#[no_mangle]
+pub extern "C" fn fst_set_streambuilder_add_le(ptr: *mut set::StreamBuilder<'static>,
+                                               c_bound: *mut libc::c_char)
+                                               -> *mut set::StreamBuilder<'static> {
+    let sb = val_from_ptr!(ptr);
+    to_raw_ptr(sb.le(cstr_to_str(c_bound)))
+}
+
 #[no_mangle]
 pub extern "C" fn fst_set_streambuilder_add_lt(ptr: *mut set::StreamBuilder<'static>,
                                                c_bound: *mut libc::c_char)

rust_fst/__init__.py

@@ -1,4 +1,4 @@
-from .set import Set
+from .set import Set, UnionSet
 from .map import Map
 
-__all__ = ["Set", "Map"]
+__all__ = ["Set", "UnionSet", "Map"]

rust_fst/set.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from enum import Enum
 
 from .common import KeyStreamIterator
 from .lib import ffi, lib, checked_call
@@ -55,14 +56,40 @@ def get_set(self):
         return Set(None, _pointer=self._set_ptr)
 
 
+class OpBuilderInputType(Enum):
+    SET = 1
+    STREAM_BUILDER = 2
+
+
 class OpBuilder(object):
-    def __init__(self, set_ptr):
+
+    _BUILDERS = {
+        OpBuilderInputType.SET: lib.fst_set_make_opbuilder,
+        OpBuilderInputType.STREAM_BUILDER: lib.fst_set_make_opbuilder_streambuilder,
+    }
+    _PUSHERS  =  {
+        OpBuilderInputType.SET: lib.fst_set_opbuilder_push,
+        OpBuilderInputType.STREAM_BUILDER: lib.fst_set_opbuilder_push_streambuilder,
+    }
+
+    @classmethod
+    def from_slice(cls, set_ptr, s):
+        sb = StreamBuilder.from_slice(set_ptr, s)
+        opbuilder = OpBuilder(sb._ptr,
+                              input_type=OpBuilderInputType.STREAM_BUILDER)
+        return opbuilder
+
+    def __init__(self, ptr, input_type=OpBuilderInputType.SET):
+        if input_type not in self._BUILDERS:
+            raise ValueError(
+                "input_type must be a member of OpBuilderInputType.")
+        self._input_type = input_type
         # NOTE: No need for `ffi.gc`, since the struct will be free'd
         #       once we call union/intersection/difference
-        self._ptr = lib.fst_set_make_opbuilder(set_ptr)
+        self._ptr = OpBuilder._BUILDERS[self._input_type](ptr)
 
-    def push(self, set_ptr):
-        lib.fst_set_opbuilder_push(self._ptr, set_ptr)
+    def push(self, ptr):
+        OpBuilder._PUSHERS[self._input_type](self._ptr, ptr)
 
     def union(self):
         stream_ptr = lib.fst_set_opbuilder_union(self._ptr)
@@ -86,6 +113,44 @@ def symmetric_difference(self):
                                  lib.fst_set_symmetricdifference_free)
 
 
+class StreamBuilder(object):
+
+    @classmethod
+    def from_slice(cls, set_ptr, slice_bounds):
+        sb = StreamBuilder(set_ptr)
+        if slice_bounds.start:
+            sb.ge(slice_bounds.start)
+        if slice_bounds.stop:
+            sb.lt(slice_bounds.stop)
+        return sb
+
+    def __init__(self, set_ptr):
+        # NOTE: No need for `ffi.gc`, since the struct will be free'd
+        #       once we call union/intersection/difference
+        self._ptr = lib.fst_set_streambuilder_new(set_ptr)
+
+    def finish(self):
+        stream_ptr = lib.fst_set_streambuilder_finish(self._ptr)
+        return KeyStreamIterator(stream_ptr, lib.fst_set_stream_next,
+                                 lib.fst_set_stream_free)
+
+    def ge(self, bound):
+        c_start = ffi.new("char[]", bound.encode('utf8'))
+        self._ptr = lib.fst_set_streambuilder_add_ge(self._ptr, c_start)
+
+    def gt(self, bound):
+        c_start = ffi.new("char[]", bound.encode('utf8'))
+        self._ptr = lib.fst_set_streambuilder_add_gt(self._ptr, c_start)
+
+    def le(self, bound):
+        c_end = ffi.new("char[]", bound.encode('utf8'))
+        self._ptr = lib.fst_set_streambuilder_add_le(self._ptr, c_end)
+
+    def lt(self, bound):
+        c_end = ffi.new("char[]", bound.encode('utf8'))
+        self._ptr = lib.fst_set_streambuilder_add_lt(self._ptr, c_end)
+
+
 class Set(object):
     """ An immutable ordered string set backed by a finite state transducer.
 
@@ -203,19 +268,11 @@ def __getitem__(self, s):
         if s.start and s.stop and s.start > s.stop:
             raise ValueError(
                 "Start key must be lexicographically smaller than stop.")
-        sb_ptr = lib.fst_set_streambuilder_new(self._ptr)
-        if s.start:
-            c_start = ffi.new("char[]", s.start.encode('utf8'))
-            sb_ptr = lib.fst_set_streambuilder_add_ge(sb_ptr, c_start)
-        if s.stop:
-            c_stop = ffi.new("char[]", s.stop.encode('utf8'))
-            sb_ptr = lib.fst_set_streambuilder_add_lt(sb_ptr, c_stop)
-        stream_ptr = lib.fst_set_streambuilder_finish(sb_ptr)
-        return KeyStreamIterator(stream_ptr, lib.fst_set_stream_next,
-                                 lib.fst_set_stream_free)
+        sb = StreamBuilder.from_slice(self._ptr, s)
+        return sb.finish()
 
     def _make_opbuilder(self, *others):
-        opbuilder = OpBuilder(self._ptr)
+        opbuilder = OpBuilder(self._ptr, input_type=OpBuilderInputType.SET)
         for oth in others:
             opbuilder.push(oth._ptr)
         return opbuilder
@@ -333,3 +390,65 @@ def search(self, term, max_dist):
         return KeyStreamIterator(stream_ptr, lib.fst_set_levstream_next,
                                  lib.fst_set_levstream_free, lev_ptr,
                                  lib.fst_levenshtein_free)
+
+
+class UnionSet(object):
+    """ A collection of Set objects that offer efficient operations across all
+    members.
+    """
+    def __init__(self, *sets):
+        self.sets = list(sets)
+
+    def __contains__(self, val):
+        """ Check if the set contains the value. """
+        return any([
+            lib.fst_set_contains(fst._ptr,
+                                 ffi.new("char[]",
+                                         val.encode('utf8')))
+            for fst in self.sets
+        ])
+
+    def __getitem__(self, s):
+        """ Get an iterator over a range of set contents.
+
+        Start and stop indices of the slice must be unicode strings.
+
+        .. important::
+            Slicing follows the semantics for numerical indices, i.e. the
+            `stop` value is **exclusive**. For example, given the set
+            `s = Set.from_iter(["bar", "baz", "foo", "moo"])`, `s['b': 'f']`
+            will only return `"bar"` and `"baz"`.
+
+        :param s:   A slice that specifies the range of the set to retrieve
+        :type s:    :py:class:`slice`
+        """
+        if not isinstance(s, slice):
+            raise ValueError(
+                "Value must be a string slice (e.g. `['foo':]`)")
+        if s.start and s.stop and s.start > s.stop:
+            raise ValueError(
+                "Start key must be lexicographically smaller than stop.")
+        if len(self.sets) <= 1:
+            raise ValueError(
+                "Must have more than one set to operate on.")
+
+        opbuilder = OpBuilder.from_slice(self.sets[0]._ptr, s)
+        streams = []
+        for fst in self.sets[1:]:
+            sb = StreamBuilder.from_slice(fst._ptr, s)
+            streams.append(sb)
+        for sb in streams:
+            opbuilder.push(sb._ptr)
+        return opbuilder.union()
+
+    def __iter__(self):
+        """ Get an iterator over all keys in all sets in lexicographical order.
+        """
+        if len(self.sets) <= 1:
+            raise ValueError(
+                "Must have more than one set to operate on.")
+        opbuilder = OpBuilder(self.sets[0]._ptr,
+                              input_type=OpBuilderInputType.SET)
+        for fst in self.sets[1:]:
+            opbuilder.push(fst._ptr)
+        return opbuilder.union()

tests/test_set.py

@@ -2,10 +2,11 @@
 import pytest
 
 import rust_fst.lib as lib
-from rust_fst import Set
+from rust_fst import Set, UnionSet
 
 
 TEST_KEYS = [u"möö", "bar", "baz", "foo"]
+TEST_KEYS2 = ["bing", "baz", "bap", "foo"]
 
 
 def do_build(path, keys=TEST_KEYS, sorted_=True):
@@ -21,6 +22,17 @@ def fst_set(tmpdir):
     return Set(fst_path)
 
 
+@pytest.fixture
+def fst_unionset(tmpdir):
+    fst_path1 = str(tmpdir.join('test1.fst'))
+    fst_path2 = str(tmpdir.join('test2.fst'))
+    do_build(fst_path1, keys=TEST_KEYS)
+    do_build(fst_path2, keys=TEST_KEYS2)
+    a = Set(fst_path1)
+    b = Set(fst_path2)
+    return UnionSet(a, b)
+
+
 def test_build(tmpdir):
     fst_path = tmpdir.join('test.fst')
     do_build(str(fst_path))
@@ -147,3 +159,23 @@ def test_range(fst_set):
         fst_set['c':'a']
     with pytest.raises(ValueError):
         fst_set['c']
+
+
+def test_unionset_contains(fst_unionset):
+    for key in TEST_KEYS+TEST_KEYS2:
+        assert key in fst_unionset
+
+
+def test_unionset_iter(fst_unionset):
+    stored_keys = list(fst_unionset)
+    assert stored_keys == sorted(set(TEST_KEYS+TEST_KEYS2))
+
+
+def test_unionset_range(fst_unionset):
+    assert list(fst_unionset['f':]) == ['foo', u'möö']
+    assert list(fst_unionset[:'m']) == ['bap', 'bar', 'baz', 'bing', 'foo']
+    assert list(fst_unionset['baz':'m']) == ['baz', 'bing', 'foo']
+    with pytest.raises(ValueError):
+        fst_unionset['c':'a']
+    with pytest.raises(ValueError):
+        fst_unionset['c']