]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
Register array dumpers for oid lookups
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sat, 28 Aug 2021 03:21:20 +0000 (05:21 +0200)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sat, 28 Aug 2021 03:21:20 +0000 (05:21 +0200)
This allows, among other things, to dump composite types containing array
in binary.

psycopg/psycopg/_adapters_map.py
psycopg/psycopg/types/array.py
tests/types/test_composite.py

index 4e301bec8b40b7662bedce14a9c2990efe78f11b..4fdcc9c0ce15d492b7ee9ff6267834f2801f6a17 100644 (file)
@@ -105,7 +105,7 @@ class AdaptersMap:
         return None
 
     def register_dumper(
-        self, cls: Union[type, str], dumper: Type[Dumper]
+        self, cls: Union[type, str, None], dumper: Type[Dumper]
     ) -> None:
         """
         Configure the context to use *dumper* to convert object of type *cls*.
@@ -123,7 +123,7 @@ class AdaptersMap:
         case it should be the fully qualified name of the object (e.g.
         ``"uuid.UUID"``).
         """
-        if not isinstance(cls, (str, type)):
+        if not (cls is None or isinstance(cls, (str, type))):
             raise TypeError(
                 f"dumpers should be registered on classes, got {cls} instead"
             )
@@ -133,12 +133,13 @@ class AdaptersMap:
 
         # Register the dumper both as its format and as auto
         # so that the last dumper registered is used in auto (%s) format
-        for fmt in (PyFormat.from_pq(dumper.format), PyFormat.AUTO):
-            if not self._own_dumpers[fmt]:
-                self._dumpers[fmt] = self._dumpers[fmt].copy()
-                self._own_dumpers[fmt] = True
+        if cls:
+            for fmt in (PyFormat.from_pq(dumper.format), PyFormat.AUTO):
+                if not self._own_dumpers[fmt]:
+                    self._dumpers[fmt] = self._dumpers[fmt].copy()
+                    self._own_dumpers[fmt] = True
 
-            self._dumpers[fmt][cls] = dumper
+                self._dumpers[fmt][cls] = dumper
 
         # Register the dumper by oid, if the oid of the dumper is fixed
         if dumper.oid:
index b58cb520a533e86d3afa5a283341d8dd6a978c07..3238a59f5b995004e6f2edbb96a71a1d5ddaaaa6 100644 (file)
@@ -35,9 +35,16 @@ TEXT_ARRAY_OID = postgres.types["text"].array_oid
 
 
 class BaseListDumper(RecursiveDumper):
+    element_oid = 0
+
     def __init__(self, cls: type, context: Optional[AdaptContext] = None):
         super().__init__(cls, context)
         self.sub_dumper: Optional[Dumper] = None
+        if self.element_oid and context:
+            sdclass = context.adapters.get_dumper_by_oid(
+                self.element_oid, self.format
+            )
+            self.sub_dumper = sdclass(type(None), context)
 
     def _find_list_element(self, L: List[Any]) -> Any:
         """
@@ -417,23 +424,45 @@ class ArrayBinaryLoader(BaseArrayLoader):
 def register_array(
     info: TypeInfo, context: Optional[AdaptContext] = None
 ) -> None:
+    if not info.array_oid:
+        raise ValueError(f"the type info {info} doesn't describe an array")
+
     adapters = context.adapters if context else postgres.adapters
 
-    base: Type[BaseArrayLoader] = ArrayLoader
-    lname = f"{info.name.title()}{base.__name__}"
+    base: Type = ArrayLoader
+    name = f"{info.name.title()}{base.__name__}"
     attribs = {
         "base_oid": info.oid,
         "delimiter": info.delimiter.encode("utf-8"),
     }
-    loader = type(lname, (base,), attribs)
+    loader = type(name, (base,), attribs)
     adapters.register_loader(info.array_oid, loader)
 
     base = ArrayBinaryLoader
-    lname = f"{info.name.title()}{base.__name__}"
+    name = f"{info.name.title()}{base.__name__}"
     attribs = {"base_oid": info.oid}
-    loader = type(lname, (base,), attribs)
+    loader = type(name, (base,), attribs)
     adapters.register_loader(info.array_oid, loader)
 
+    base = ListDumper
+    name = f"{info.name.title()}{base.__name__}"
+    attribs = {
+        "oid": info.array_oid,
+        "element_oid": info.oid,
+        "delimiter": info.delimiter.encode("utf-8"),
+    }
+    dumper = type(name, (base,), attribs)
+    adapters.register_dumper(None, dumper)
+
+    base = ListBinaryDumper
+    name = f"{info.name.title()}{base.__name__}"
+    attribs = {
+        "oid": info.array_oid,
+        "element_oid": info.oid,
+    }
+    dumper = type(name, (base,), attribs)
+    adapters.register_dumper(None, dumper)
+
 
 def register_default_adapters(context: AdaptContext) -> None:
     # The text dumper is more flexible as it can handle lists of mixed type,
index 8b0e97d5c0de7c1b296cf04040a8b34ee2491db7..2c314e9cfc076c05ed8337f4c9acd49a5ff1657c 100644 (file)
@@ -4,6 +4,7 @@ from psycopg import pq, postgres
 from psycopg.sql import Identifier
 from psycopg.adapt import PyFormat as Format
 from psycopg.postgres import types as builtins
+from psycopg.types.range import Range
 from psycopg.types.composite import CompositeInfo, register_composite
 from psycopg.types.composite import TupleDumper, TupleBinaryDumper
 
@@ -63,6 +64,25 @@ def test_load_all_chars(conn, fmt_out):
     assert res == (s,)
 
 
+@pytest.mark.parametrize("fmt_in", [Format.AUTO, Format.TEXT, Format.BINARY])
+def test_dump_builtin_empty_range(conn, fmt_in):
+    conn.execute(
+        """
+        drop type if exists tmptype;
+        create type tmptype as (num integer, range daterange, nums integer[])
+        """
+    )
+    info = CompositeInfo.fetch(conn, "tmptype")
+    register_composite(info, conn)
+
+    cur = conn.execute(
+        f"select pg_typeof(%{fmt_in})",
+        [info.python_type(10, Range(empty=True), [])],
+    )
+    print(cur._query.params[0])
+    assert cur.fetchone()[0] == "tmptype"
+
+
 @pytest.mark.parametrize(
     "rec, want",
     [