diff --git a/docs/python-api.rst b/docs/python-api.rst index fae8cee..2ded7b0 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -1415,3 +1415,25 @@ You can use it in code like this: if spatialite: db.conn.enable_load_extension(True) db.conn.load_extension(spatialite) + +.. _python_api_register_function: + +Registering custom SQL functions +================================ + +SQLite supports registering custom SQL functions written in Python. The ``@db.register_function`` decorator provides syntactic sugar for registering these functions. + +.. code-block:: python + + from sqlite_utils import Database + + db = Database(memory=True) + + @db.register_function + def reverse_string(s): + return "".join(reversed(list(s))) + + print(db.execute('select reverse_string("hello")').fetchone())[0] + # This prints "olleh" + +If your Python function accepts multiple arguments the SQL version will accept the same number of arguments. diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index ebffabf..daf3e25 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -4,6 +4,7 @@ import contextlib import datetime import decimal import hashlib +import inspect import itertools import json import os @@ -143,6 +144,11 @@ class Database: def __repr__(self): return "".format(self.conn) + def register_function(self, fn): + name = fn.__name__ + arity = len(inspect.signature(fn).parameters) + self.conn.create_function(name, arity, fn) + def execute(self, sql, parameters=None): if self._tracer: self._tracer(sql, parameters) diff --git a/tests/test_register_function.py b/tests/test_register_function.py new file mode 100644 index 0000000..4aa5a3b --- /dev/null +++ b/tests/test_register_function.py @@ -0,0 +1,16 @@ +def test_register_function(fresh_db): + @fresh_db.register_function + def reverse_string(s): + return "".join(reversed(list(s))) + + result = fresh_db.execute('select reverse_string("hello")').fetchone()[0] + assert result == "olleh" + + +def test_register_function_multiple_arguments(fresh_db): + @fresh_db.register_function + def a_times_b_plus_c(a, b, c): + return a * b + c + + result = fresh_db.execute("select a_times_b_plus_c(2, 3, 4)").fetchone()[0] + assert result == 10