Add support for regex search to UnionSet.

Author: davidblewett

rust/rust_fst.h

@@ -66,6 +66,7 @@ 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_regexstream(SetRegexStream*);
 SetOpBuilder* fst_set_make_opbuilder_union(SetUnion*);
 void fst_set_free(Set*);
 
@@ -80,6 +81,7 @@ 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_regexstream(SetOpBuilder*, SetRegexStream*);
 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

@@ -154,6 +154,14 @@ pub extern "C" fn fst_set_make_opbuilder_levstream(ptr: *mut SetLevStream) -> *m
     to_raw_ptr(ob)
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_make_opbuilder_regexstream(ptr: *mut SetRegexStream) -> *mut set::OpBuilder<'static> {
+    let srs = val_from_ptr!(ptr);
+    let mut ob = set::OpBuilder::new();
+    ob.push(srs.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);
@@ -184,6 +192,13 @@ pub extern "C" fn fst_set_opbuilder_push_levstream(ptr: *mut set::OpBuilder<'sta
     ob.push(sls.into_stream());
 }
 
+#[no_mangle]
+pub extern "C" fn fst_set_opbuilder_push_regexstream(ptr: *mut set::OpBuilder<'static>, srs_ptr: *mut SetRegexStream) {
+    let srs = val_from_ptr!(srs_ptr);
+    let ob = mutref_from_ptr!(ptr);
+    ob.push(srs.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

@@ -61,6 +61,7 @@ class OpBuilderInputType(Enum):
     STREAM_BUILDER = 2
     UNION = 3
     SEARCH = 4
+    SEARCH_RE = 5
 
 
 def _build_levsearch(fst, term, max_dist):
@@ -72,19 +73,28 @@ def _build_levsearch(fst, term, max_dist):
     return lib.fst_set_levsearch(fst._ptr, lev_ptr)
 
 
+def _build_research(fst, pattern):
+    re_ptr = checked_call(
+        lib.fst_regex_new, fst._ctx,
+        ffi.new("char[]", pattern.encode('utf8')))
+    return lib.fst_set_regexsearch(fst._ptr, re_ptr)
+
+
 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,
         OpBuilderInputType.SEARCH: lib.fst_set_make_opbuilder_levstream,
+        OpBuilderInputType.SEARCH_RE: lib.fst_set_make_opbuilder_regexstream,
     }
     _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,
+        OpBuilderInputType.SEARCH_RE: lib.fst_set_opbuilder_push_regexstream,
     }
 
     @classmethod
@@ -94,6 +104,13 @@ def from_search(cls, fst, term, max_dist):
                               input_type=OpBuilderInputType.SEARCH)
         return opbuilder
 
+    @classmethod
+    def from_search_re(cls, fst, pattern):
+        stream_ptr = _build_research(fst, pattern)
+        opbuilder = OpBuilder(stream_ptr,
+                              input_type=OpBuilderInputType.SEARCH_RE)
+        return opbuilder
+
     @classmethod
     def from_slice(cls, set_ptr, s):
         sb = StreamBuilder.from_slice(set_ptr, s)
@@ -535,6 +552,38 @@ def search(self, term, max_dist):
             opbuilder.push(_build_levsearch(fst, term, max_dist))
         return opbuilder.union()
 
+    def search_re(self, pattern):
+        """ Search the set with a regular expression.
+
+        Note that the regular expression syntax is not Python's, but the one
+        supported by the `regex` Rust crate, which is almost identical
+        to the engine of the RE2 engine.
+
+        For a documentation of the syntax, see:
+        http://doc.rust-lang.org/regex/regex/index.html#syntax
+
+        Due to limitations of the underlying FST, only a subset of this syntax
+        is supported. Most notably absent are:
+
+        * Lazy quantifiers (``r'*?'``, ``r'+?'``)
+        * Word boundaries (``r'\\b'``)
+        * Other zero-width assertions (``r'^'``, ``r'$'``)
+
+        For background on these limitations, consult the documentation of
+        the Rust crate: http://burntsushi.net/rustdoc/fst/struct.Regex.html
+
+        :param pattern:     A regular expression
+        :returns:           An iterator over all matching keys 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_re(self.sets[0], pattern)
+        for fst in self.sets[1:]:
+            opbuilder.push(_build_research(fst, pattern))
+        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

@@ -200,6 +200,11 @@ def test_unionset_search(fst_unionset):
     assert matches == ["bap", "bar", "baz"]
 
 
+def test_unionset_search_re(fst_unionset):
+    matches = list(fst_unionset.search_re(r'ba.*'))
+    assert matches == ["bap", "bar", "baz"]
+
+
 def test_unionset_symmetric_difference():
     a = Set.from_iter(["bar", "foo"])
     b = Set.from_iter(["baz", "foo"])