Add support for Levenshtein fuzzy search to UnionSet.

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_levstream(SetLevStream*);
 SetOpBuilder* fst_set_make_opbuilder_union(SetUnion*);
 void fst_set_free(Set*);
 
@@ -78,6 +79,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_levstream(SetOpBuilder*, SetLevStream*);
 void fst_set_opbuilder_push_streambuilder(SetOpBuilder*, SetStreamBuilder*);
 void fst_set_opbuilder_push_union(SetOpBuilder*, SetUnion*);
 void fst_set_opbuilder_free(SetOpBuilder*);

rust/src/set.rs

@@ -146,6 +146,14 @@ 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_levstream(ptr: *mut SetLevStream) -> *mut set::OpBuilder<'static> {
+    let sls = val_from_ptr!(ptr);
+    let mut ob = set::OpBuilder::new();
+    ob.push(sls.into_stream());
+    to_raw_ptr(ob)
+}
+
 #[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);
@@ -169,6 +177,13 @@ pub extern "C" fn fst_set_opbuilder_push(ptr: *mut set::OpBuilder, set_ptr: *mut
     ob.push(set);
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_opbuilder_push_levstream(ptr: *mut set::OpBuilder<'static>, sls_ptr: *mut SetLevStream) {
+    let sls = val_from_ptr!(sls_ptr);
+    let ob = mutref_from_ptr!(ptr);
+    ob.push(sls.into_stream());
+}
+
 #[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);

rust_fst/set.py

@@ -60,6 +60,16 @@ class OpBuilderInputType(Enum):
     SET = 1
     STREAM_BUILDER = 2
     UNION = 3
+    SEARCH = 4
+
+
+def _build_levsearch(fst, term, max_dist):
+    lev_ptr = checked_call(
+        lib.fst_levenshtein_new,
+        fst._ctx,
+        ffi.new("char[]", term.encode('utf8')),
+        max_dist)
+    return lib.fst_set_levsearch(fst._ptr, lev_ptr)
 
 
 class OpBuilder(object):
@@ -68,13 +78,22 @@ class OpBuilder(object):
         OpBuilderInputType.SET: lib.fst_set_make_opbuilder,
         OpBuilderInputType.STREAM_BUILDER: lib.fst_set_make_opbuilder_streambuilder,
         OpBuilderInputType.UNION: lib.fst_set_make_opbuilder_union,
+        OpBuilderInputType.SEARCH: lib.fst_set_make_opbuilder_levstream,
     }
     _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,
+        OpBuilderInputType.SEARCH: lib.fst_set_opbuilder_push_levstream,
     }
 
+    @classmethod
+    def from_search(cls, fst, term, max_dist):
+        stream_ptr = _build_levsearch(fst, term, max_dist)
+        opbuilder = OpBuilder(stream_ptr,
+                              input_type=OpBuilderInputType.SEARCH)
+        return opbuilder
+
     @classmethod
     def from_slice(cls, set_ptr, s):
         sb = StreamBuilder.from_slice(set_ptr, s)
@@ -500,6 +519,22 @@ def intersection(self, *others):
         """
         return self._make_opbuilder(*others).intersection()
 
+    def search(self, term, max_dist):
+        """ Search the set with a Levenshtein automaton.
+
+        :param term:        The search term
+        :param max_dist:    The maximum edit distance for search results
+        :returns:           Iterator over matching values in the set
+        :rtype:             :py:class:`KeyStreamIterator`
+        """
+        if len(self.sets) <= 1:
+            raise ValueError(
+                "Must have more than one set to operate on.")
+        opbuilder = OpBuilder.from_search(self.sets[0], term, max_dist)
+        for fst in self.sets[1:]:
+            opbuilder.push(_build_levsearch(fst, term, max_dist))
+        return opbuilder.union()
+
     def symmetric_difference(self, *others):
         """ Get an iterator over the keys in the symmetric difference of this
             set and others.

tests/test_set.py

@@ -195,6 +195,11 @@ def test_unionset_range(fst_unionset):
         fst_unionset['c']
 
 
+def test_unionset_search(fst_unionset):
+    matches = list(fst_unionset.search("bam", 1))
+    assert matches == ["bap", "bar", "baz"]
+
+
 def test_unionset_symmetric_difference():
     a = Set.from_iter(["bar", "foo"])
     b = Set.from_iter(["baz", "foo"])