From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 26 Aug 2021 15:24:26 +0000 (+0200)
Subject: Register composite dumper if the factory is a type
X-Git-Tag: 3.0.beta1~30^2~4
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=01655037bb42244ba9dbbc08111f2e359689df4a;p=thirdparty%2Fpsycopg.git

Register composite dumper if the factory is a type

Expose the type as info.python_type after registering.
---

diff --git a/psycopg/psycopg/_typeinfo.py b/psycopg/psycopg/_typeinfo.py
index 17f572fda..0f7ff2026 100644
--- a/psycopg/psycopg/_typeinfo.py
+++ b/psycopg/psycopg/_typeinfo.py
@@ -191,6 +191,8 @@ class CompositeInfo(TypeInfo):
         super().__init__(name, oid, array_oid)
         self.field_names = field_names
         self.field_types = field_types
+        # Will be set by register() if the `factory` is a type
+        self.python_type: Optional[type] = None
 
     def register(
         self,
diff --git a/psycopg/psycopg/types/composite.py b/psycopg/psycopg/types/composite.py
index 711238400..e1a843946 100644
--- a/psycopg/psycopg/types/composite.py
+++ b/psycopg/psycopg/types/composite.py
@@ -212,6 +212,14 @@ def register_composite(
     )
     adapters.register_loader(info.oid, loader)
 
+    # If the factory is a type, register a dumper for it
+    if isinstance(factory, type):
+        dumper = type(
+            f"{info.name.title()}Dumper", (TupleDumper,), {"_oid": info.oid}
+        )
+        adapters.register_dumper(factory, dumper)
+        info.python_type = factory
+
 
 def register_default_adapters(context: AdaptContext) -> None:
     adapters = context.adapters
diff --git a/tests/types/test_composite.py b/tests/types/test_composite.py
index 749c7376b..952959d93 100644
--- a/tests/types/test_composite.py
+++ b/tests/types/test_composite.py
@@ -4,7 +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.composite import CompositeInfo
+from psycopg.types.composite import CompositeInfo, TupleDumper
 
 
 tests_str = [
@@ -198,6 +198,7 @@ def test_load_composite_factory(conn, testcomp, fmt_out):
             self.foo, self.bar, self.baz = args
 
     info.register(conn, factory=MyThing)
+    assert info.python_type is MyThing
 
     cur = conn.cursor(binary=fmt_out)
     res = cur.execute("select row('hello', 10, 20)::testcomp").fetchone()[0]
@@ -220,6 +221,11 @@ def test_register_scope(conn, testcomp):
         for oid in (info.oid, info.array_oid):
             assert postgres.adapters._loaders[fmt].pop(oid)
 
+    for fmt in (Format.AUTO, Format.TEXT):
+        assert postgres.adapters._dumpers[fmt].pop(info.python_type)
+
+    assert info.python_type not in postgres.adapters._dumpers[Format.BINARY]
+
     cur = conn.cursor()
     info.register(cur)
     for fmt in (pq.Format.TEXT, pq.Format.BINARY):
@@ -233,3 +239,31 @@ def test_register_scope(conn, testcomp):
         for oid in (info.oid, info.array_oid):
             assert oid not in postgres.adapters._loaders[fmt]
             assert oid in conn.adapters._loaders[fmt]
+
+
+def test_type_dumper_registered(conn, testcomp):
+    info = CompositeInfo.fetch(conn, "testcomp")
+    info.register(conn)
+    assert issubclass(info.python_type, tuple)
+    assert info.python_type.__name__ == "testcomp"
+    d = conn.adapters.get_dumper(info.python_type, "s")
+    assert issubclass(d, TupleDumper)
+    assert d is not TupleDumper
+
+    tc = info.python_type("foo", 42, 3.14)
+    cur = conn.execute("select pg_typeof(%s)", [tc])
+    assert cur.fetchone()[0] == "testcomp"
+
+
+def test_callable_dumper_not_registered(conn, testcomp):
+    info = CompositeInfo.fetch(conn, "testcomp")
+
+    def fac(*args):
+        return args + (args[-1],)
+
+    info.register(conn, factory=fac)
+    assert info.python_type is None
+
+    # but the loader is registered
+    cur = conn.execute("select '(foo,42,3.14)'::testcomp")
+    assert cur.fetchone()[0] == ("foo", 42, 3.14, 3.14)