feat: allow binary JSON dumps

Some JSON libraries, in particular the widespread
[orjson](https://github.com/ijl/orjson#quickstart) library, directly
dump data in binary format.

This change allows it and avoids having to decode/encode in these cases
by testing the return type of the `dumps` function.

diff --git a/docs/news.rst b/docs/news.rst
index d213bfcacfb1f76a5255f02b5e0f8507b7396d7a..38f5e037c800e144d198fd4fbf4bb1367b53d11c 100644 (file)
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -24,6 +24,8 @@ Psycopg 3.1.9
 - Fix loading ROW values with different types in the same query using the
   binary protocol (:ticket:`#545`).
 - Fix dumping recursive composite types (:ticket:`#547`).
+- Allow JSON dumpers to dump `bytes` directly instead of `str`,
+  for better compatibility with libraries like orjson and msgspec (:ticket:`#569`)
 
 
 Psycopg 3.1.8

diff --git a/psycopg/psycopg/types/json.py b/psycopg/psycopg/types/json.py
index b3323cf0c2de71f12e1ff9ad5f531dc219c7ad8e..702215bdc832008d6a633a7237d07b8ca37280ce 100644 (file)
--- a/psycopg/psycopg/types/json.py
+++ b/psycopg/psycopg/types/json.py
@@ -14,7 +14,7 @@ from ..pq import Format
 from ..adapt import Buffer, Dumper, Loader, PyFormat, AdaptersMap
 from ..errors import DataError
 
-JsonDumpsFunction = Callable[[Any], str]
+JsonDumpsFunction = Callable[[Any], Union[str, bytes]]
 JsonLoadsFunction = Callable[[Union[str, bytes]], Any]
 
 
@@ -132,7 +132,10 @@ class _JsonDumper(Dumper):
             obj = obj.obj
         else:
             dumps = self.dumps
-        return dumps(obj).encode()
+        data = dumps(obj)
+        if isinstance(data, str):
+            return data.encode()
+        return data
 
 
 class JsonDumper(_JsonDumper):

diff --git a/tests/types/test_json.py b/tests/types/test_json.py
index ba476adb7ecd0e16d7b38c9c75920a4a3bb72d38..3fce60b4ca8132272f0cb8244802a7ad84aef48f 100644 (file)
--- a/tests/types/test_json.py
+++ b/tests/types/test_json.py
@@ -141,6 +141,21 @@ def test_dump_customise(conn, wrapper, fmt_in):
         set_json_dumps(json.dumps)
 
 
+@pytest.mark.parametrize("fmt_in", PyFormat)
+@pytest.mark.parametrize("wrapper", ["Json", "Jsonb"])
+def test_dump_customise_bytes(conn, wrapper, fmt_in):
+    wrapper = getattr(psycopg.types.json, wrapper)
+    obj = {"foo": "bar"}
+    cur = conn.cursor()
+
+    set_json_dumps(my_dumps_bytes)
+    try:
+        cur.execute(f"select %{fmt_in.value}->>'baz' = 'qux'", (wrapper(obj),))
+        assert cur.fetchone()[0] is True
+    finally:
+        set_json_dumps(json.dumps)
+
+
 @pytest.mark.parametrize("fmt_in", PyFormat)
 @pytest.mark.parametrize("wrapper", ["Json", "Jsonb"])
 def test_dump_customise_context(conn, wrapper, fmt_in):
@@ -205,6 +220,12 @@ def my_dumps(obj):
     return json.dumps(obj)
 
 
+def my_dumps_bytes(obj):
+    obj = deepcopy(obj)
+    obj["baz"] = "qux"
+    return json.dumps(obj).encode()
+
+
 def my_loads(data):
     obj = json.loads(data)
     obj["answer"] = 42