register_function(name=...) argument, closes #458

This commit is contained in:
Simon Willison 2022-07-27 17:13:49 -07:00
commit 1491b66dd7
3 changed files with 30 additions and 6 deletions

View file

@ -2513,6 +2513,16 @@ You can also use the method as a function decorator like so:
print(db.execute('select reverse_string("hello")').fetchone()[0])
By default, the name of the Python function will be used as the name of the SQL function. You can customize this with the ``name=`` keyword argument:
.. code-block:: python
@db.register_function(name="rev")
def reverse_string(s):
return "".join(reversed(list(s)))
print(db.execute('select rev("hello")').fetchone()[0])
Python 3.8 added the ability to register `deterministic SQLite functions <https://sqlite.org/deterministic.html>`__, allowing you to indicate that a function will return the exact same result for any given inputs and hence allowing SQLite to apply some performance optimizations. You can mark a function as deterministic using ``deterministic=True``, like this:
.. code-block:: python

View file

@ -367,7 +367,11 @@ class Database:
return "<Database {}>".format(self.conn)
def register_function(
self, fn: Callable = None, deterministic: bool = False, replace: bool = False
self,
fn: Callable = None,
deterministic: bool = False,
replace: bool = False,
name: Optional[str] = None,
):
"""
``fn`` will be made available as a function within SQL, with the same name and number
@ -388,12 +392,13 @@ class Database:
:param fn: Function to register
:param deterministic: set ``True`` for functions that always returns the same output for a given input
:param replace: set ``True`` to replace an existing function with the same name - otherwise throw an error
:param name: name of the SQLite function - if not specified, the Python function name will be used
"""
def register(fn):
name = fn.__name__
fn_name = name or fn.__name__
arity = len(inspect.signature(fn).parameters)
if not replace and (name, arity) in self._registered_functions:
if not replace and (fn_name, arity) in self._registered_functions:
return fn
kwargs = {}
registered = False
@ -401,15 +406,15 @@ class Database:
# Try this, but fall back if sqlite3.NotSupportedError
try:
self.conn.create_function(
name, arity, fn, **dict(kwargs, deterministic=True)
fn_name, arity, fn, **dict(kwargs, deterministic=True)
)
registered = True
except (sqlite3.NotSupportedError, TypeError):
# TypeError is Python 3.7 "function takes at most 3 arguments"
pass
if not registered:
self.conn.create_function(name, arity, fn, **kwargs)
self._registered_functions.add((name, arity))
self.conn.create_function(fn_name, arity, fn, **kwargs)
self._registered_functions.add((fn_name, arity))
return fn
if fn is None:

View file

@ -14,6 +14,15 @@ def test_register_function(fresh_db):
assert result == "olleh"
def test_register_function_custom_name(fresh_db):
@fresh_db.register_function(name="revstr")
def reverse_string(s):
return "".join(reversed(list(s)))
result = fresh_db.execute('select revstr("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):