@db.register_function(..., replace=True), closes #199

This commit is contained in:
Simon Willison 2020-11-06 07:53:22 -08:00
commit 27b67f1cae
4 changed files with 38 additions and 2 deletions

View file

@ -1719,3 +1719,13 @@ Python 3.8 added the ability to register `deterministic SQLite functions <https:
return "".join(reversed(list(s)))
If you run this on a version of Python prior to 3.8 your code will still work, but the ``deterministic=True`` parameter will be ignored.
By default registering a function with the same name and number of arguments will have no effect - the ``Database`` instance keeps track of functions that have already been registered and skips registering them if ``@db.register_function`` is called a second time.
If you want to deliberately replace the registered function with a new implementation, use the ``replace=True`` argument:
.. code-block:: python
@db.register_function(deterministic=True, replace=True)
def reverse_string(s):
return s[::-1]

View file

@ -22,7 +22,7 @@ setup(
version=VERSION,
license="Apache License, Version 2.0",
packages=find_packages(exclude=["tests", "tests.*"]),
install_requires=["click", "click-default-group", "tabulate"],
install_requires=["sqlite-fts4", "click", "click-default-group", "tabulate"],
setup_requires=["pytest-runner"],
extras_require={
"test": ["pytest", "black", "hypothesis"],

View file

@ -154,6 +154,7 @@ class Database:
self._tracer = tracer
if recursive_triggers:
self.execute("PRAGMA recursive_triggers=on;")
self._registered_functions = set()
@contextlib.contextmanager
def tracer(self, tracer=None):
@ -170,14 +171,17 @@ class Database:
def __repr__(self):
return "<Database {}>".format(self.conn)
def register_function(self, fn=None, deterministic=None):
def register_function(self, fn=None, deterministic=None, replace=False):
def register(fn):
name = fn.__name__
arity = len(inspect.signature(fn).parameters)
if not replace and (name, arity) in self._registered_functions:
return fn
kwargs = {}
if deterministic and sys.version_info >= (3, 8):
kwargs["deterministic"] = True
self.conn.create_function(name, arity, fn, **kwargs)
self._registered_functions.add((name, arity))
return fn
if fn is None:

View file

@ -44,3 +44,25 @@ def test_register_function_deterministic_registered(fresh_db):
fresh_db.conn.create_function.assert_called_with(
"to_lower_2", 1, to_lower_2, deterministic=True
)
def test_register_function_replace(fresh_db):
@fresh_db.register_function()
def one():
return "one"
assert "one" == fresh_db.execute("select one()").fetchone()[0]
# This will fail to replace the function:
@fresh_db.register_function()
def one():
return "two"
assert "one" == fresh_db.execute("select one()").fetchone()[0]
# This will replace it
@fresh_db.register_function(replace=True)
def one():
return "two"
assert "two" == fresh_db.execute("select one()").fetchone()[0]