Add Database(json_default=, json_object_hook=) for custom JSON encoding/decoding

Values not natively serializable (set, enum.Enum, project-specific types)
previously always fell back to repr() when written to a JSON column, with
no way to reconstruct them on read. json_default is passed through to
json.dumps() as default= wherever the library serializes dict/list/tuple
values; json_object_hook is passed through to json.loads() as object_hook=
when deserialize_json=True. Both are Database-level, matching the existing
deserialize_json option, since every write/read path already threads
through self.db.

Fixes #521

docs/python-api.rst

@@ -2302,6 +2302,56 @@ any string value that looks like a JSON object or array back into a ``dict`` or
 This defaults to ``False``, so existing code that expects raw strings continues to work
 unchanged.
 
+.. _python_api_custom_json:
+
+Custom JSON encoding and decoding
+----------------------------------
+
+By default, values that are not natively JSON-serializable - such as
+``set``, ``enum.Enum`` members or custom classes - are converted using
+``repr()`` when they are written to a JSON column, since ``dict``, ``list``
+and ``tuple`` values are always serialized with ``json.dumps()``. To take
+full control of that conversion, pass a ``json_default`` function to the
+``Database()`` constructor - it is used as the ``default=`` callback to
+``json.dumps()``:
+
+.. code-block:: python
+
+    import enum
+
+    class Color(enum.Enum):
+        RED = "red"
+        BLUE = "blue"
+
+    def encode(value):
+        if isinstance(value, set):
+            return list(value)
+        if isinstance(value, enum.Enum):
+            return value.value
+        raise TypeError(f"Cannot serialize {value!r}")
+
+    db = sqlite_utils.Database("data.db", json_default=encode)
+    db["examples"].insert({"id": 1, "tags": {"a", "b"}, "color": Color.RED})
+
+To reconstruct custom types when reading rows back with
+``deserialize_json=True``, pass a ``json_object_hook`` function - it is
+used as the ``object_hook=`` callback to ``json.loads()`` and is called
+with every JSON object (``dict``) encountered while parsing:
+
+.. code-block:: python
+
+    def decode(obj):
+        if "__color__" in obj:
+            return Color(obj["__color__"])
+        return obj
+
+    db = sqlite_utils.Database(
+        "data.db", deserialize_json=True, json_object_hook=decode
+    )
+
+Both options apply to every table on that ``Database`` connection - there is
+no per-table equivalent, matching how ``deserialize_json`` works.
+
 .. _python_api_conversions:
 
 Converting column values using SQL functions

sqlite_utils/db.py

@@ -505,6 +505,13 @@ class Database:
       look like JSON objects or arrays back into ``dict``/``list`` objects when rows are read
       via ``.rows``, ``.rows_where()``, ``.get()``, ``.search()`` or ``.query()``. Defaults to
       ``False``, so values are returned as the raw strings stored in the database.
+    :param json_default: optional function used as the ``default=`` callback to ``json.dumps()``
+      when serializing values (dicts, lists, tuples) to JSON on write. Use this to support types
+      that are not natively JSON-serializable, such as ``set``, ``enum.Enum`` or custom classes.
+      Defaults to ``repr``.
+    :param json_object_hook: optional function used as the ``object_hook=`` callback to
+      ``json.loads()`` when ``deserialize_json=True``. Use this to reconstruct custom types from
+      the JSON objects produced by ``json_default``.
     """
 
     _counts_table_name = "_counts"
@@ -524,11 +531,15 @@ class Database:
         use_old_upsert: bool = False,
         strict: bool = False,
         deserialize_json: bool = False,
+        json_default: Callable[[Any], Any] | None = None,
+        json_object_hook: Callable[[dict], Any] | None = None,
     ):
         self.memory_name = None
         self.memory = False
         self.use_old_upsert = use_old_upsert
         self.deserialize_json = deserialize_json
+        self.json_default = json_default
+        self.json_object_hook = json_object_hook
         if not (
             (filename_or_conn is not None and (not memory and not memory_name))
             or (filename_or_conn is None and (memory or memory_name))
@@ -822,7 +833,10 @@ class Database:
         "Build a row dict, deserializing JSON string values if self.deserialize_json is set."
         d = dict(zip(keys, row))
         if self.deserialize_json:
-            d = {key: dejsonify_if_needed(value) for key, value in d.items()}
+            d = {
+                key: dejsonify_if_needed(value, self.json_object_hook)
+                for key, value in d.items()
+            }
         return d
 
     def query(
@@ -3761,7 +3775,7 @@ class Table(Queryable):
             sets.append(
                 "{} = {}".format(quote_identifier(key), conversions.get(key, "?"))
             )
-            args.append(jsonify_if_needed(value))
+            args.append(jsonify_if_needed(value, self.db.json_default))
         wheres = [f"{quote_identifier(pk_name)} = ?" for pk_name in pks]
         args.extend(pk_values)
         sql = "update {} set {sets} where {wheres}".format(
@@ -3841,7 +3855,7 @@ class Table(Queryable):
 
             def convert_value(v):
                 bar.update(1)
-                return jsonify_if_needed(fn(v))
+                return jsonify_if_needed(fn(v), self.db.json_default)
 
             fn_name = getattr(fn, "__name__", "fn")
             if fn_name == "<lambda>":
@@ -3968,11 +3982,14 @@ class Table(Queryable):
                 # Pad short records with None, truncate long ones
                 record_len = len(record)
                 if record_len < num_columns:
-                    record_values = [jsonify_if_needed(v) for v in record] + [None] * (
-                        num_columns - record_len
-                    )
+                    record_values = [
+                        jsonify_if_needed(v, self.db.json_default) for v in record
+                    ] + [None] * (num_columns - record_len)
                 else:
-                    record_values = [jsonify_if_needed(v) for v in record[:num_columns]]
+                    record_values = [
+                        jsonify_if_needed(v, self.db.json_default)
+                        for v in record[:num_columns]
+                    ]
                 # Only process extracts if there are any
                 if has_extracts:
                     for i, key in enumerate(all_columns):
@@ -3994,7 +4011,8 @@ class Table(Queryable):
                                 if key != hash_id
                                 else hash_record(record, hash_id_columns)
                             ),
-                        )
+                        ),
+                        self.db.json_default,
                     )
                     if key in extracts and value is not None:
                         extract_table = extracts[key]
@@ -5105,11 +5123,13 @@ class View(Queryable):
                 raise
 
 
-def jsonify_if_needed(value: object) -> object:
+def jsonify_if_needed(
+    value: object, json_default: Callable[[Any], Any] | None = None
+) -> object:
     if isinstance(value, decimal.Decimal):
         return float(value)
     if isinstance(value, (dict, list, tuple)):
-        return json.dumps(value, default=repr, ensure_ascii=False)
+        return json.dumps(value, default=json_default or repr, ensure_ascii=False)
     elif isinstance(value, (datetime.time, datetime.date, datetime.datetime)):
         return value.isoformat()
     elif isinstance(value, (datetime.timedelta, uuid.UUID)):
@@ -5118,14 +5138,16 @@ def jsonify_if_needed(value: object) -> object:
         return value
 
 
-def dejsonify_if_needed(value: object) -> object:
+def dejsonify_if_needed(
+    value: object, json_object_hook: Callable[[dict], Any] | None = None
+) -> object:
     "Parse a string that looks like a JSON object or array back into a dict/list."
     if isinstance(value, str) and value[:1] in ("{", "["):
         try:
-            parsed = json.loads(value)
+            parsed = json.loads(value, object_hook=json_object_hook)
         except ValueError:
             return value
-        if isinstance(parsed, (dict, list)):
+        if isinstance(parsed, (dict, list)) or json_object_hook is not None:
             return parsed
     return value
 

tests/test_create.py

@@ -921,6 +921,46 @@ def test_deserialize_json(fresh_db):
     assert list(db2.query("select data from test"))[0]["data"] == {"foo": "bar"}
 
 
+def test_json_default_and_object_hook(fresh_db):
+    import enum
+
+    class Color(enum.Enum):
+        RED = "red"
+
+    def encode(value):
+        if isinstance(value, set):
+            return sorted(value)
+        if isinstance(value, enum.Enum):
+            return {"__enum__": value.value}
+        raise TypeError(value)
+
+    def decode(obj):
+        if "__enum__" in obj:
+            return Color(obj["__enum__"])
+        return obj
+
+    # A set/enum can't be a column's own type (SQLite has no such type), but they
+    # are common inside a dict/list column value - that's what json_default is for.
+    db = Database(fresh_db.conn, json_default=encode)
+    db["test"].insert(
+        {"id": 1, "data": {"tags": {"b", "a"}, "color": Color.RED}},
+        pk="id",
+    )
+    row = db.execute("select data from test").fetchone()
+    assert row[0] == '{"tags": ["a", "b"], "color": {"__enum__": "red"}}'
+
+    # Without json_default, an unsupported nested type falls back to repr()
+    db_default = Database(fresh_db.conn)
+    db_default["test2"].insert({"id": 1, "data": {"tags": {"a"}}}, pk="id")
+    row2 = db_default.execute("select data from test2").fetchone()
+    assert row2[0] == json.dumps({"tags": repr({"a"})})
+
+    # json_object_hook reconstructs the custom type on read
+    db2 = Database(fresh_db.conn, deserialize_json=True, json_object_hook=decode)
+    row3 = db2["test"].get(1)
+    assert row3["data"] == {"tags": ["a", "b"], "color": Color.RED}
+
+
 def test_insert_list_nested_unicode(fresh_db):
     fresh_db["test"].insert(
         {"id": 1, "data": {"key1": {"nested": ["cømplex"]}}}, pk="id"