From 0789bad8f7581fd96dec5bde51a75e937dffb1e0 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 28 Oct 2020 14:24:03 -0700
Subject: [PATCH] @db.register_function(deterministic=True), closes #191

---
 docs/python-api.rst             | 10 ++++++++++
 sqlite_utils/db.py              | 19 +++++++++++++++----
 tests/test_register_function.py | 30 ++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index cb7461b..79ad510 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -1703,3 +1703,13 @@ SQLite supports registering custom SQL functions written in Python. The ``@db.re
     # This prints "olleh"
 
 If your Python function accepts multiple arguments the SQL version will accept the same number of arguments.
+
+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
+
+    @db.register_function(deterministic=True)
+    def reverse_string(s):
+        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.
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 877c17f..2e70f9d 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -10,6 +10,7 @@ import itertools
 import json
 import os
 import pathlib
+import sys
 import textwrap
 import uuid
 
@@ -149,10 +150,20 @@ class Database:
     def __repr__(self):
         return "<Database {}>".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 register_function(self, fn=None, deterministic=None):
+        def register(fn):
+            name = fn.__name__
+            arity = len(inspect.signature(fn).parameters)
+            kwargs = {}
+            if deterministic and sys.version_info >= (3, 8):
+                kwargs["deterministic"] = True
+            self.conn.create_function(name, arity, fn, **kwargs)
+            return fn
+
+        if fn is None:
+            return register
+        else:
+            register(fn)
 
     def execute(self, sql, parameters=None):
         if self._tracer:
diff --git a/tests/test_register_function.py b/tests/test_register_function.py
index 4aa5a3b..dab0a8d 100644
--- a/tests/test_register_function.py
+++ b/tests/test_register_function.py
@@ -1,3 +1,8 @@
+import pytest
+import sys
+from unittest.mock import MagicMock
+
+
 def test_register_function(fresh_db):
     @fresh_db.register_function
     def reverse_string(s):
@@ -14,3 +19,28 @@ def test_register_function_multiple_arguments(fresh_db):
 
     result = fresh_db.execute("select a_times_b_plus_c(2, 3, 4)").fetchone()[0]
     assert result == 10
+
+
+def test_register_function_deterministic(fresh_db):
+    @fresh_db.register_function(deterministic=True)
+    def to_lower(s):
+        return s.lower()
+
+    result = fresh_db.execute("select to_lower('BOB')").fetchone()[0]
+    assert result == "bob"
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 8), reason="deterministic=True was added in Python 3.8"
+)
+def test_register_function_deterministic_registered(fresh_db):
+    fresh_db.conn = MagicMock()
+    fresh_db.conn.create_function = MagicMock()
+
+    @fresh_db.register_function(deterministic=True)
+    def to_lower_2(s):
+        return s.lower()
+
+    fresh_db.conn.create_function.assert_called_with(
+        "to_lower_2", 1, to_lower_2, deterministic=True
+    )