add 'fts5' module to write FTS5 tokenizer

see #5

Author: hideaki-t

sqlitefts/__init__.py

@@ -1,5 +1,4 @@
-from .tokenizer import register_tokenizer
-from .fts3 import Tokenizer, make_tokenizer_module
+from .fts3 import Tokenizer, make_tokenizer_module, register_tokenizer
 from . import tokenizer, ranking
 
 __all__ = ["Tokenizer", "make_tokenizer_module", "register_tokenizer",

sqlitefts/fts3.py

@@ -3,8 +3,13 @@
 PoC SQLite FTS5 tokenizer in Python
 """
 from __future__ import print_function, unicode_literals
+import sys
+import struct
 
-from .tokenizer import ffi, SQLITE_OK, SQLITE_DONE
+from .tokenizer import (ffi, dll, get_db_from_connection, SQLITE_OK,
+                        SQLITE_DONE)
+
+SQLITE_DBCONFIG_ENABLE_FTS3_TOKENIZER = 1004
 
 ffi.cdef('''
 typedef struct sqlite3_tokenizer_module sqlite3_tokenizer_module;
@@ -127,3 +132,21 @@ def xclose(pCursor):
     tokenizer_modules[tokenizer] = (tokenizer_module, xcreate, xdestroy, xopen,
                                     xclose, xnext)
     return tokenizer_module
+
+
+def enable_fts3_tokenizer(c):
+    db = get_db_from_connection(c)
+    rc = dll.sqlite3_db_config(db, SQLITE_DBCONFIG_ENABLE_FTS3_TOKENIZER,
+                               ffi.cast('int', 1), ffi.NULL)
+    return rc == 0
+
+
+def register_tokenizer(c, name, tokenizer_module):
+    """ register tokenizer module with SQLite connection. """
+    module_addr = int(ffi.cast('uintptr_t', tokenizer_module))
+    address_blob = struct.pack("P", module_addr)
+    if sys.version_info.major == 2:
+        address_blob = buffer(address_blob)
+    enable_fts3_tokenizer(c)
+    r = c.execute('SELECT fts3_tokenizer(?, ?)', (name, address_blob))
+    return r

sqlitefts/tokenizer.py

@@ -3,20 +3,13 @@
 a proof of concept implementation of SQLite FTS tokenizers in Python
 """
 from __future__ import print_function, unicode_literals
-
 import sys
-import struct
 
 from cffi import FFI
 
 SQLITE_OK = 0
 SQLITE_DONE = 101
 
-if sys.version_info.major == 2:
-    global buffer
-else:
-    buffer = lambda x: x
-
 ffi = FFI()
 ffi.cdef('''
 typedef struct sqlite3 sqlite3;
@@ -43,31 +36,11 @@
     dll = ffi.dlopen(find_library("sqlite3"))
 
 
-def f():
-    SQLITE_DBCONFIG_ENABLE_FTS3_TOKENIZER = 1004
-
-    def enable_fts3_tokenizer(c):
-        db = getattr(c, '_db', None)
-        if db:
-            # pypy's SQLite3 connection has _db using cffi
-            db = ffi.cast('sqlite3*', db)
-        else:
-            db = ffi.cast('PyObject *', id(c)).db
-        rc = dll.sqlite3_db_config(db, SQLITE_DBCONFIG_ENABLE_FTS3_TOKENIZER,
-                                   ffi.cast('int', 1), ffi.NULL)
-        return rc == 0
-
-    return enable_fts3_tokenizer
-
-
-enable_fts3_tokenizer = f()
-del f
-
-
-def register_tokenizer(c, name, tokenizer_module):
-    """ register tokenizer module with SQLite connection. """
-    module_addr = int(ffi.cast('uintptr_t', tokenizer_module))
-    address_blob = buffer(struct.pack("P", module_addr))
-    enable_fts3_tokenizer(c)
-    r = c.execute('SELECT fts3_tokenizer(?, ?)', (name, address_blob))
-    return r
+def get_db_from_connection(c):
+    db = getattr(c, '_db', None)
+    if db:
+        # pypy's SQLite3 connection has _db using cffi
+        db = ffi.cast('sqlite3*', db)
+    else:
+        db = ffi.cast('PyObject *', id(c)).db
+    return db

tests/test_fts5.py

@@ -0,0 +1,260 @@
+# coding: utf-8
+from __future__ import print_function, unicode_literals
+import sqlite3
+import re
+
+from cffi import FFI
+
+from sqlitefts import fts5
+from sqlitefts import Tokenizer
+
+ffi = FFI()
+
+
+class SimpleTokenizer(Tokenizer):
+    _p = re.compile(r'\w+', re.UNICODE)
+
+    def tokenize(self, text):
+        for m in self._p.finditer(text):
+            s, e = m.span()
+            t = text[s:e]
+            l = len(t.encode('utf-8'))
+            p = len(text[:s].encode('utf-8'))
+            yield t, p, p + l
+
+
+def test_fts5_api_from_db():
+    c = sqlite3.connect(':memory:')
+    fts5api = fts5.fts5_api_from_db(c)
+    print(fts5api)
+    assert fts5api.iVersion == 2
+    assert fts5api.xCreateTokenizer
+    c.close()
+
+
+def test_make_tokenizer():
+    c = sqlite3.connect(':memory:')
+    tm = fts5.make_fts5_tokenizer(SimpleTokenizer())
+    assert all(
+        getattr(tm, x) is not None
+        for x in ('xCreate', 'xDelete', 'xTokenize'))
+    c.close()
+
+
+def test_register_tokenizer():
+    name = 'simpe'
+    c = sqlite3.connect(':memory:')
+    tm = fts5.make_fts5_tokenizer(SimpleTokenizer())
+    assert fts5.register_tokenizer(c, name, tm)
+    c.close()
+
+
+def test_register_tokenizer_with_destroy():
+    name = 'simpe'
+    c = sqlite3.connect(':memory:')
+
+    arg_on_destroy = None
+    context = "hello"
+
+    def on_destroy(x):
+        nonlocal arg_on_destroy
+        arg_on_destroy = x
+
+    tm = fts5.make_fts5_tokenizer(SimpleTokenizer())
+    assert fts5.register_tokenizer(
+        c, name, tm, context=context, on_destroy=on_destroy)
+    c.close()
+    assert arg_on_destroy == context
+
+
+def test_createtable():
+    c = sqlite3.connect(':memory:')
+    c.row_factory = sqlite3.Row
+    name = 'super_simple'
+    sql = "CREATE VIRTUAL TABLE fts USING fts5(w, tokenize={})".format(name)
+    fts5.register_tokenizer(c, name,
+                            fts5.make_fts5_tokenizer(SimpleTokenizer()))
+    c.execute(sql)
+
+    r = c.execute(
+        "SELECT * FROM sqlite_master WHERE type='table' AND name='fts'").fetchone(
+        )
+    assert r
+    assert r[str('type')] == 'table' and r[str('name')] == 'fts' and r[str(
+        'tbl_name')] == 'fts'
+    assert r[str('sql')].upper() == sql.upper()
+    c.close()
+
+
+def test_insert():
+    c = sqlite3.connect(':memory:')
+    c.row_factory = sqlite3.Row
+    name = 'super_simple'
+    content = 'これは日本語で書かれています'
+    fts5.register_tokenizer(c, name,
+                            fts5.make_fts5_tokenizer(SimpleTokenizer()))
+    c.execute(
+        "CREATE VIRTUAL TABLE fts USING FTS5(content, tokenize={})".format(
+            name))
+    r = c.execute('INSERT INTO fts VALUES(?)', (content, ))
+    assert r.rowcount == 1
+    r = c.execute("SELECT * FROM fts").fetchone()
+    assert r
+    assert r[str('content')] == content
+    c.close()
+
+
+def test_match():
+    c = sqlite3.connect(':memory:')
+    c.row_factory = sqlite3.Row
+    name = 'super_simple'
+    contents = [('abc def', ), ('abc xyz', ), ('あいうえお かきくけこ', ),
+                ('あいうえお らりるれろ', )]
+    fts5.register_tokenizer(c, name,
+                            fts5.make_fts5_tokenizer(SimpleTokenizer()))
+    c.execute(
+        "CREATE VIRTUAL TABLE fts USING FTS5(content, tokenize={})".format(
+            name))
+    r = c.executemany('INSERT INTO fts VALUES(?)', contents)
+    assert r.rowcount == 4
+    r = c.execute("SELECT * FROM fts").fetchall()
+    assert len(r) == 4
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'abc'").fetchall()
+    assert len(r) == 2
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'def'").fetchall()
+    assert len(r) == 1 and r[0][str('content')] == contents[0][0]
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'xyz'").fetchall()
+    assert len(r) == 1 and r[0][str('content')] == contents[1][0]
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'zzz'").fetchall()
+    assert len(r) == 0
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'あいうえお'").fetchall()
+    assert len(r) == 2
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'かきくけこ'").fetchall()
+    assert len(r) == 1 and r[0][str('content')] == contents[2][0]
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'らりるれろ'").fetchall()
+    assert len(r) == 1 and r[0][str('content')] == contents[3][0]
+    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'まみむめも'").fetchall()
+    assert len(r) == 0
+    c.close()
+
+
+def test_full_text_index_queries():
+    name = 'super_simple'
+    docs = [(
+        'README',
+        'sqlitefts-python provides binding for tokenizer of SQLite Full-Text search(FTS3/4). It allows you to write tokenizers in Python.'
+    ), ('LICENSE',
+        '''Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:'''),
+            ('日本語', 'あいうえお かきくけこ さしすせそ たちつてと なにぬねの')]
+    with sqlite3.connect(':memory:') as c:
+        c.row_factory = sqlite3.Row
+        fts5.register_tokenizer(c, name,
+                                fts5.make_fts5_tokenizer(SimpleTokenizer()))
+        c.execute(
+            "CREATE VIRTUAL TABLE docs USING FTS5(title, body, tokenize={})".format(
+                name))
+        c.executemany("INSERT INTO docs(title, body) VALUES(?, ?)", docs)
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Python'").fetchall(
+        )
+        assert len(r) == 1
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'bind'").fetchall()
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'binding'").fetchall()
+        assert len(r) == 1
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'to'").fetchall()
+        assert len(r) == 2
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお'").fetchall()
+        assert len(r) == 1
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'らりるれろ'").fetchall()
+        assert len(r) == 0
+        assert (
+            c.execute(
+                "SELECT * FROM docs WHERE docs MATCH 'binding'").fetchall()[0]
+            == c.execute(
+                "SELECT * FROM docs WHERE docs MATCH 'body:binding'").fetchall(
+                )[0])
+        assert (c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'body:binding'").fetchall(
+            )[0] == c.execute(
+                "SELECT * FROM docs WHERE docs MATCH 'body:binding'").fetchall(
+                )[0])
+        assert (
+            c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお'").fetchall(
+            )[0] == c.execute(
+                "SELECT * FROM docs WHERE docs MATCH 'body:あいうえお'").fetchall()[
+                    0])
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'title:bind'").fetchall()
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'title:README'").fetchall()
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'title:日本語'").fetchall()
+        assert len(r) == 1
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'to in'").fetchall()
+        assert len(r) == 2
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Py*'").fetchall()
+        assert len(r) == 1
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Z*'").fetchall()
+        assert len(r) == 0
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あ*'").fetchall()
+        assert len(r) == 1
+        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'ん*'").fetchall()
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'tokenizer SQLite'").fetchall(
+            )
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH '\"tokenizer SQLite\"'").fetchall(
+            )
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'あいうえお たちつてと'").fetchall()
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH '\"あいうえお たちつてと\"'").fetchall()
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'tok* + SQL*'").fetchall()
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'tok* of SQL*'").fetchall()
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'あ* + さ*'").fetchall()
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'あ* かきくけこ さ*'").fetchall()
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'NEAR(tokenizer SQLite)'").fetchall(
+            )
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'NEAR(binding SQLite, 2)'").fetchall(
+            )
+        assert len(r) == 0
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'NEAR(binding SQLite, 3)'").fetchall(
+            )
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'NEAR(あいうえお たちつてと)'").fetchall(
+            )
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'NEAR(あいうえお たちつてと, 2)'").fetchall(
+            )
+        assert len(r) == 1
+        r = c.execute(
+            "SELECT * FROM docs WHERE docs MATCH 'NEAR(あいうえお たちつてと, 3)'").fetchall(
+            )
+        assert len(r) == 1