From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Mon, 8 Feb 2021 22:56:22 +0000 (+0100)
Subject: Added set_jsonb_loads/dumps functions to customise JSON adaptation
X-Git-Tag: 3.0.dev0~120
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=07824b87c7dd2a85a189d5a3b5bd364877c1a90d;p=thirdparty%2Fpsycopg.git

Added set_jsonb_loads/dumps functions to customise JSON adaptation

Dropped dumps function from the Json/Jsonb wrappers.

Added JSON documentation.
---

diff --git a/docs/adapt-types.rst b/docs/adapt-types.rst
index adcd23f5d..21da06dac 100644
--- a/docs/adapt-types.rst
+++ b/docs/adapt-types.rst
@@ -249,12 +249,65 @@ time and more bandwidth. See :ref:`binary-data` for details.
 .. __: https://www.postgresql.org/docs/current/datatype-binary.html
 
 
+.. _adapt-json:
+
+JSON adaptation
+---------------
+
+`!psycopg3` can map between Python objects and PostgreSQL `json/jsonb
+types`__, allowing to customise the load and dump function used.
+
+.. __: https://www.postgresql.org/docs/current/datatype-json.html
+
+Because several Python objects could be considered JSON (dicts, lists,
+scalars, even date/time if using a dumps function customised to use them),
+`!psycopg3` requires you to wrap what you want to dump as JSON into a wrapper:
+either `psycogp3.types.Json` or `~psycopg3.types.Jsonb`.
+
+.. code:: python
+
+    from psycopg3.types import Jsonb
+
+    thing = {"foo": ["bar", 42]}
+    conn.execute("insert into mytable values (%s)", [Jsonb(thing)])
+
+By default `!psycopg3` uses the standard library `json.dumps()`__ and
+`json.loads()`__ functions to serialize and de-serialize Python objects to
+JSON. If you want to customise globally how serialization happens, for
+instance changing serialization parameters or using a different JSON library,
+you can specify your own functions using the `psycopg3.types.set_json_dumps()`
+and `~psycopg3.types.set_json_loads()` functions.
+
+..
+    weird: intersphinx doesn't work
+
+.. __: https://docs.python.org/3/library/json.html#json.dumps
+.. __: https://docs.python.org/3/library/json.html#json.loads
+
+.. code:: python
+
+    from functools import partial
+    from psycopg3.types import Jsonb, set_json_dumps, set_json_loads
+    import ujson
+
+    # Use a faster dump function
+    set_json_dumps(ujson.dumps)
+
+    # Return floating point values as Decimal
+    set_json_loads(partial(json.loads, parse_float=Decimal))
+
+    conn.execute("select %s", [Jsonb({"value": 123.45})]).fetchone()[0]
+    # {'value': Decimal('123.45')}
+
+If you need a more precise customisation, such as per-connection instead of
+global, you can subclass and register the JSON adapters in the right context:
+see :ref:`json-adapters`.
+
 .. _adapt-date:
 .. _adapt-list:
 .. _adapt-composite:
 .. _adapt-hstore:
 .. _adapt-range:
-.. _adapt-json:
 .. _adapt-uuid:
 .. _adapt-network:
 
@@ -264,6 +317,3 @@ TODO adaptation
 .. admonition:: TODO
 
     Document the other types
-
-    Document that empty array don't roundtrip in text mode and require
-    a cast in binary to be used in any context.
diff --git a/docs/types.rst b/docs/types.rst
index 3f144bbb8..40ff77213 100644
--- a/docs/types.rst
+++ b/docs/types.rst
@@ -105,3 +105,54 @@ Objects wrappers
     - Int2, Int4, Int8, ...
     - Json, Jsonb
     - Range
+
+
+.. _json-adapters:
+
+JSON adapters
+-------------
+
+.. autoclass:: Json
+.. autoclass:: Jsonb
+
+Wrappers to signal to convert *obj* to a json or jsonb PostgreSQL value.
+
+Any object supported by the underlying `!dumps()` function can be wrapped.
+
+
+.. autofunction:: set_json_dumps
+.. autofunction:: set_json_loads
+
+.. autoclass:: JsonDumper
+
+    .. automethod:: get_dumps
+
+.. autoclass:: JsonBinaryDumper
+.. autoclass:: JsonbDumper
+.. autoclass:: JsonbBinaryDumper
+
+`~psycopg3.adapt.Dumper` subclasses using the function provided by
+`set_json_dumps()` function to serialize the Python object wrapped by
+`Json`/`Jsonb`.
+
+If you need to specify different `!dumps()` functions in different contexts
+you can subclass one/some of these functions to override the
+`~JsonDumper.get_dumps()` method and `~psycopg3.adapt.Dumper.register()` them
+on the right connection or cursor.
+
+.. autoclass:: JsonLoader
+
+    .. automethod:: get_loads
+
+.. autoclass:: JsonBinaryLoader
+.. autoclass:: JsonbLoader
+.. autoclass:: JsonbBinaryLoader
+
+`~psycopg3.adapt.Loader` subclasses using the function provided by
+`set_json_loads()` function to de-serialize :sql:`json`/:sql:`jsonb`
+PostgreSQL values to Python objects.
+
+If you need to specify different `!loads()` functions in different contexts
+you can subclass one/some of these functions to override the
+`~JsonLoader.get_loads()` method and `~psycopg3.adapt.Loader.register()` them
+on the right connection or cursor.
diff --git a/psycopg3/psycopg3/types/__init__.py b/psycopg3/psycopg3/types/__init__.py
index 18d7e5247..2f211714f 100644
--- a/psycopg3/psycopg3/types/__init__.py
+++ b/psycopg3/psycopg3/types/__init__.py
@@ -18,6 +18,9 @@ from .range import Range
 # Database types descriptors
 from .._typeinfo import TypeInfo, RangeInfo, CompositeInfo
 
+# Json global registrations
+from .json import set_json_dumps, set_json_loads
+
 # Adapter objects
 from .text import (
     StringDumper,
@@ -82,6 +85,7 @@ from .json import (
     JsonbBinaryDumper,
     JsonLoader,
     JsonBinaryLoader,
+    JsonbLoader,
     JsonbBinaryLoader,
 )
 from .uuid import (
@@ -192,7 +196,7 @@ def register_default_globals(ctx: AdaptContext) -> None:
     JsonbDumper.register(Jsonb, ctx)
     JsonbBinaryDumper.register(Jsonb, ctx)
     JsonLoader.register("json", ctx)
-    JsonLoader.register("jsonb", ctx)
+    JsonbLoader.register("jsonb", ctx)
     JsonBinaryLoader.register("json", ctx)
     JsonbBinaryLoader.register("jsonb", ctx)
 
diff --git a/psycopg3/psycopg3/types/json.py b/psycopg3/psycopg3/types/json.py
index 1e021cabc..528fc78b0 100644
--- a/psycopg3/psycopg3/types/json.py
+++ b/psycopg3/psycopg3/types/json.py
@@ -5,22 +5,57 @@ Adapers for JSON types.
 # Copyright (C) 2020-2021 The Psycopg Team
 
 import json
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Union
 
 from ..pq import Format
 from ..oids import postgres_types as builtins
 from ..adapt import Buffer, Dumper, Loader
+from ..proto import AdaptContext
 from ..errors import DataError
 
 JsonDumpsFunction = Callable[[Any], str]
+JsonLoadsFunction = Callable[[Union[str, bytes, bytearray]], Any]
+
+# Global load/dump functions, used by default.
+_loads: JsonLoadsFunction = json.loads
+_dumps: JsonDumpsFunction = json.dumps
+
+
+def set_json_dumps(dumps: JsonDumpsFunction) -> None:
+    """
+    Set a global JSON serialisation function to use by default by JSON dumpers.
+
+    Defaults to the builtin `json.dumps()`. You can override it to use a
+    different JSON library or to use customised arguments.
+
+    If you need a non-global customisation you can subclass the `!JsonDumper`
+    family of classes, overriding the `!get_loads()` method, and register
+    your class in the context required.
+    """
+    global _dumps
+    _dumps = dumps
+
+
+def set_json_loads(loads: JsonLoadsFunction) -> None:
+    """
+    Set a global JSON parsing function to use by default by the JSON loaders.
+
+    Defaults to the builtin `json.loads()`. You can override it to use a
+    different JSON library or to use customised arguments.
+
+    If you need a non-global customisation you can subclass the `!JsonLoader`
+    family of classes, overriding the `!get_loads()` method, and register
+    your class in the context required.
+    """
+    global _loads
+    _loads = loads
 
 
 class _JsonWrapper:
-    __slots__ = ("obj", "_dumps")
+    __slots__ = ("obj",)
 
-    def __init__(self, obj: Any, dumps: Optional[JsonDumpsFunction] = None):
+    def __init__(self, obj: Any):
         self.obj = obj
-        self._dumps: JsonDumpsFunction = dumps or json.dumps
 
     def __repr__(self) -> str:
         sobj = repr(self.obj)
@@ -28,9 +63,6 @@ class _JsonWrapper:
             sobj = f"{sobj[:35]} ... ({len(sobj)} chars)"
         return f"{self.__class__.__name__}({sobj})"
 
-    def dumps(self) -> str:
-        return self._dumps(self.obj)
-
 
 class Json(_JsonWrapper):
     __slots__ = ()
@@ -44,8 +76,21 @@ class _JsonDumper(Dumper):
 
     format = Format.TEXT
 
+    def __init__(self, cls: type, context: Optional[AdaptContext] = None):
+        super().__init__(cls, context)
+        self._dumps = self.get_dumps()
+
+    def get_dumps(self) -> JsonDumpsFunction:
+        r"""
+        Return a `json.dumps()`\-compatible function to serialize the object.
+
+        Subclasses can override this function to specify custom JSON
+        serialization per context.
+        """
+        return _dumps
+
     def dump(self, obj: _JsonWrapper) -> bytes:
-        return obj.dumps().encode("utf-8")
+        return self._dumps(obj.obj).encode("utf-8")
 
 
 class JsonDumper(_JsonDumper):
@@ -70,26 +115,43 @@ class JsonbBinaryDumper(JsonbDumper):
     format = Format.BINARY
 
     def dump(self, obj: _JsonWrapper) -> bytes:
-        return b"\x01" + obj.dumps().encode("utf-8")
+        return b"\x01" + self._dumps(obj.obj).encode("utf-8")
 
 
-class JsonLoader(Loader):
+class _JsonLoader(Loader):
+    def __init__(self, oid: int, context: Optional[AdaptContext] = None):
+        super().__init__(oid, context)
+        self._loads = self.get_loads()
 
-    format = Format.TEXT
+    def get_loads(self) -> JsonLoadsFunction:
+        r"""
+        Return a `json.loads()`\-compatible function to de-serialize the value.
+
+        Subclasses can override this function to specify custom JSON
+        de-serialization per context.
+        """
+        return _loads
 
     def load(self, data: Buffer) -> Any:
-        # Json crashes on memoryview
+        # json.loads() cannot work on memoryview.
         if isinstance(data, memoryview):
             data = bytes(data)
-        return json.loads(data)
+        return self._loads(data)
 
 
-class JsonBinaryLoader(JsonLoader):
+class JsonLoader(_JsonLoader):
+    format = Format.TEXT
+
+
+class JsonbLoader(_JsonLoader):
+    format = Format.TEXT
+
 
+class JsonBinaryLoader(_JsonLoader):
     format = Format.BINARY
 
 
-class JsonbBinaryLoader(Loader):
+class JsonbBinaryLoader(_JsonLoader):
 
     format = Format.BINARY
 
@@ -99,4 +161,4 @@ class JsonbBinaryLoader(Loader):
         data = data[1:]
         if isinstance(data, memoryview):
             data = bytes(data)
-        return json.loads(data)
+        return self._loads(data)
diff --git a/tests/types/test_json.py b/tests/types/test_json.py
index bea17f3da..a903e8517 100644
--- a/tests/types/test_json.py
+++ b/tests/types/test_json.py
@@ -1,4 +1,5 @@
 import json
+from copy import deepcopy
 
 import pytest
 
@@ -7,6 +8,7 @@ from psycopg3 import pq
 from psycopg3 import sql
 from psycopg3.types import Json, Jsonb
 from psycopg3.adapt import Format
+from psycopg3.types import set_json_dumps, set_json_loads
 
 samples = [
     "null",
@@ -70,27 +72,78 @@ def test_json_dump_customise(conn, wrapper, fmt_in):
     wrapper = getattr(psycopg3.types, wrapper)
     obj = {"foo": "bar"}
     cur = conn.cursor()
-    cur.execute(
-        f"select %{fmt_in}->>'baz' = 'qux'", (wrapper(obj, dumps=my_dumps),)
-    )
-    assert cur.fetchone()[0] is True
+
+    set_json_dumps(my_dumps)
+    try:
+        cur.execute(f"select %{fmt_in}->>'baz' = 'qux'", (wrapper(obj),))
+        assert cur.fetchone()[0] is True
+    finally:
+        set_json_dumps(json.dumps)
 
 
 @pytest.mark.parametrize("fmt_in", [Format.AUTO, Format.TEXT, Format.BINARY])
 @pytest.mark.parametrize("wrapper", ["Json", "Jsonb"])
 def test_json_dump_subclass(conn, wrapper, fmt_in):
+    JDumper = getattr(
+        psycopg3.types,
+        f"{wrapper}{'Binary' if fmt_in != Format.TEXT else ''}Dumper",
+    )
     wrapper = getattr(psycopg3.types, wrapper)
 
-    class MyWrapper(wrapper):
-        def dumps(self):
-            return my_dumps(self.obj)
+    class MyJsonDumper(JDumper):
+        def get_dumps(self):
+            return my_dumps
 
     obj = {"foo": "bar"}
     cur = conn.cursor()
-    cur.execute(f"select %{fmt_in}->>'baz' = 'qux'", (MyWrapper(obj),))
+    MyJsonDumper.register(wrapper, context=cur)
+    cur.execute(f"select %{fmt_in}->>'baz' = 'qux'", (wrapper(obj),))
     assert cur.fetchone()[0] is True
 
 
+@pytest.mark.parametrize("binary", [True, False])
+@pytest.mark.parametrize("pgtype", ["json", "jsonb"])
+def test_json_load_customise(conn, binary, pgtype):
+    obj = {"foo": "bar"}
+    cur = conn.cursor(binary=binary)
+
+    set_json_loads(my_loads)
+    try:
+        cur.execute(f"""select '{{"foo": "bar"}}'::{pgtype}""")
+        obj = cur.fetchone()[0]
+        assert obj["foo"] == "bar"
+        assert obj["answer"] == 42
+    finally:
+        set_json_loads(json.loads)
+
+
+@pytest.mark.parametrize("binary", [True, False])
+@pytest.mark.parametrize("pgtype", ["json", "jsonb"])
+def test_json_load_subclass(conn, binary, pgtype):
+    JLoader = getattr(
+        psycopg3.types,
+        f"{pgtype.title()}{'Binary' if binary else ''}Loader",
+    )
+
+    class MyJsonLoader(JLoader):
+        def get_loads(self):
+            return my_loads
+
+    cur = conn.cursor(binary=binary)
+    MyJsonLoader.register(cur.adapters.types[pgtype].oid, context=cur)
+    cur.execute(f"""select '{{"foo": "bar"}}'::{pgtype}""")
+    obj = cur.fetchone()[0]
+    assert obj["foo"] == "bar"
+    assert obj["answer"] == 42
+
+
 def my_dumps(obj):
+    obj = deepcopy(obj)
     obj["baz"] = "qux"
     return json.dumps(obj)
+
+
+def my_loads(data):
+    obj = json.loads(data)
+    obj["answer"] = 42
+    return obj