From: Lionel Panhaleux <lionel.panhaleux@gmail.com>
Date: Fri, 19 May 2023 12:03:29 +0000 (+0200)
Subject: feat: allow binary JSON dumps
X-Git-Tag: 3.1.10~17
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=075439976abd9e4f02d18adf7cb2685fbed0118e;p=thirdparty%2Fpsycopg.git

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 d213bfcac..38f5e037c 100644
--- 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 b3323cf0c..702215bdc 100644
--- 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 ba476adb7..3fce60b4c 100644
--- 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