Added mapping from python Loader to C and context extraction

diff --git a/psycopg3/adapt.pxd b/psycopg3/adapt.pxd
index 627101acee88625b77587461d30056aecddc7c0c..4d372a3367d686a86936b943af32fe49dc8c5fb0 100644 (file)
--- a/psycopg3/adapt.pxd
+++ b/psycopg3/adapt.pxd
@@ -1,9 +1,7 @@
 from cpython.object cimport PyObject
 
 ctypedef object (*cloader_func)(const char *data, size_t length, void *context)
-# ctypedef void * (*get_context_func)(PGconn_ptr conn)
-
-cdef void register_c_loader(object pyloader, cloader_func cloader)
+ctypedef void * (*get_context_func)(object conn)
 
 
 cdef struct RowLoader:

diff --git a/psycopg3/adapt.py b/psycopg3/adapt.py
index 33cd0af792fdc8b32c5d678ef3f4e1bad3ce6252..2af9a6263f85052ec6ae2566545335f53bec2687 100644 (file)
--- a/psycopg3/adapt.py
+++ b/psycopg3/adapt.py
@@ -390,7 +390,7 @@ class UnknownLoader(Loader):
 
 
 @Loader.binary(INVALID_OID)
-def load_unknown(data: bytes) -> bytes:
+def load_unknown_binary(data: bytes) -> bytes:
     return data
 
 

diff --git a/psycopg3/adapt.pyx b/psycopg3/adapt.pyx
index aa35c8d915e341634afa3a9e88128fe8c64365ed..02cb4e1a5697bae11e897d9cf51cdf0adb5651a8 100644 (file)
--- a/psycopg3/adapt.pyx
+++ b/psycopg3/adapt.pyx
@@ -1,6 +1,6 @@
-from psycopg3.adapt cimport cloader_func
+from psycopg3.adapt cimport cloader_func, get_context_func, RowLoader
 
-from psycopg3.types cimport numeric
+from psycopg3.types cimport numeric, text
 
 
 import logging
@@ -10,11 +10,16 @@ logger = logging.getLogger("psycopg3.adapt")
 cdef class CLoader:
     cdef object pyloader
     cdef cloader_func cloader
+    cdef get_context_func get_context
 
 
 cloaders = {}
 
-cdef void register_c_loader(object pyloader, cloader_func cloader):
+cdef void register_c_loader(
+    object pyloader,
+    cloader_func cloader,
+    get_context_func get_context = NULL
+):
     """
     Map a python loader to an optimised C version.
 
@@ -24,9 +29,25 @@ cdef void register_c_loader(object pyloader, cloader_func cloader):
     cdef CLoader cl = CLoader()
     cl.pyloader = pyloader
     cl.cloader = cloader
+    cl.get_context = get_context
     cloaders[pyloader] = cl
 
 
+cdef void fill_row_loader(RowLoader *loader, object pyloader):
+    loader.pyloader = <PyObject *>pyloader
+
+    print(pyloader)
+    cdef CLoader cloader
+    cloader = cloaders.get(pyloader)
+    if cloader is not None:
+        loader.cloader = cloader.cloader
+    else:
+        cloader = cloaders.get(getattr(pyloader, '__func__', None))
+        if cloader is not None and cloader.get_context is not NULL:
+            loader.cloader = cloader.cloader
+            loader.context = cloader.get_context(pyloader.__self__)
+
+
 def register_builtin_c_loaders():
     if cloaders:
         logger.warning("c loaders already registered")
@@ -34,6 +55,7 @@ def register_builtin_c_loaders():
 
     logger.debug("registering optimised c loaders")
     register_numeric_c_loaders()
+    register_text_c_loaders()
 
 
 cdef void register_numeric_c_loaders():
@@ -44,3 +66,13 @@ cdef void register_numeric_c_loaders():
     register_c_loader(numeric.load_int4_binary, load_int4_binary)
     register_c_loader(numeric.load_int8_binary, load_int8_binary)
     register_c_loader(numeric.load_oid_binary, load_oid_binary)
+    register_c_loader(numeric.load_bool_binary, load_bool_binary)
+
+
+cdef void register_text_c_loaders():
+    logger.debug("registering optimised text c loaders")
+    from psycopg3 import adapt
+    from psycopg3.types import text
+    register_c_loader(text.StringLoader.load, load_text, get_context_text)
+    register_c_loader(adapt.UnknownLoader.load, load_unknown_text)
+    register_c_loader(adapt.load_unknown_binary, load_unknown_binary)

diff --git a/psycopg3/transform.pyx b/psycopg3/transform.pyx
index 8e77f787ca1b007c4d166fba43bd7a27f4ae8c67..2c565fdcb07cc9ebe1738aafdc1aef1009c0261a 100644 (file)
--- a/psycopg3/transform.pyx
+++ b/psycopg3/transform.pyx
@@ -157,33 +157,8 @@ cdef class Transformer:
     cdef void _set_loader(self, int col, libpq.Oid oid, int fmt):
         cdef RowLoader *loader = self._row_loaders + col
 
-        pyloader = self.get_load_function(oid, fmt)
-        self._pyloaders[col] = pyloader
-        loader.pyloader = <PyObject *>pyloader
-
-        # STUB: special-case a few loaders
-        from psycopg3.types import builtins
-
-        if oid == builtins['int2'].oid and fmt == 0:
-            loader.cloader = load_int_text
-        elif oid == builtins['int2'].oid and fmt == 1:
-            loader.cloader = load_int2_binary
-        elif oid == builtins['int4'].oid and fmt == 1:
-            loader.cloader = load_int4_binary
-        elif oid == builtins['int8'].oid and fmt == 1:
-            loader.cloader = load_int8_binary
-        elif oid == builtins['oid'].oid and fmt == 1:
-            loader.cloader = load_oid_binary
-        elif oid == builtins['bool'].oid and fmt == 1:
-            loader.cloader = load_bool_binary
-        elif oid == builtins['text'].oid:
-            loader.cloader = load_text_binary
-        elif oid == builtins['"char"'].oid:
-            loader.cloader = load_text_binary
-        elif oid == builtins['name'].oid:
-            loader.cloader = load_text_binary
-        elif oid == builtins['unknown'].oid:
-            loader.cloader = load_unknown_binary
+        pyloader = self._pyloaders[col] = self.get_load_function(oid, fmt)
+        fill_row_loader(loader, pyloader)
 
     def dump_sequence(
         self, objs: Iterable[Any], formats: Iterable[Format]

diff --git a/psycopg3/types/composite.py b/psycopg3/types/composite.py
index cf1ed0838ed9f8ebeef6fedba7ec441a8b6c0936..c801e575ec90f717a8d3045d013bec857d6245d2 100644 (file)
--- a/psycopg3/types/composite.py
+++ b/psycopg3/types/composite.py
@@ -237,7 +237,7 @@ class BinaryRecordLoader(BaseCompositeLoader):
 
     def _config_types(self, data: bytes) -> None:
         self._tx.set_row_types(
-            (oid, Format.BINARY) for oid, _, _ in self._walk_record(data)
+            [(oid, Format.BINARY) for oid, _, _ in self._walk_record(data)]
         )
 
 

diff --git a/psycopg3/types/text.pxd b/psycopg3/types/text.pxd
new file mode 100644 (file)
index 0000000..81b0425
--- /dev/null
+++ b/psycopg3/types/text.pxd
@@ -0,0 +1,4 @@
+cdef object load_text(const char *data, size_t length, void *context)
+cdef void *get_context_text(object loader)
+cdef object load_unknown_text(const char *data, size_t length, void *context)
+cdef object load_unknown_binary(const char *data, size_t length, void *context)

diff --git a/psycopg3/types/text.pyx b/psycopg3/types/text.pyx
index 1fb2fb2b427b7e270ae8e156e17de1fee3a8313f..b143d4e4eb5da5268b5d0334c9004f7218623913 100644 (file)
--- a/psycopg3/types/text.pyx
+++ b/psycopg3/types/text.pyx
@@ -1,10 +1,29 @@
+from cpython.bytes cimport PyBytes_FromStringAndSize
 from cpython.unicode cimport PyUnicode_DecodeUTF8
 
-cdef object load_text_binary(const char *data, size_t length, void *context):
+
+cdef object load_text(const char *data, size_t length, void *context):
+    # TODO: optimize
+    b = PyBytes_FromStringAndSize(data, length)
+    if context is not NULL:
+        codec = <object>context
+        # TODO: check if the refcount is right (but a DECREF here segfaults)
+        return codec(b)[0]
+    else:
+        return b
+
+
+cdef void *get_context_text(object loader):
+    if loader.decode is not None:
+        return <void *>loader.decode
+    else:
+        return NULL
+
+
+cdef object load_unknown_text(const char *data, size_t length, void *context):
     # TODO: codec
     return PyUnicode_DecodeUTF8(data, length, NULL)
 
 
 cdef object load_unknown_binary(const char *data, size_t length, void *context):
-    # TODO: codec
-    return PyUnicode_DecodeUTF8(data, length, NULL)
+    return PyBytes_FromStringAndSize(data, length)