From 05f264367260465e9e9889ca8e5e1aad017e404c Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Sun, 10 May 2020 21:05:54 +1200 Subject: [PATCH] Added mapping from python Loader to C and context extraction --- psycopg3/adapt.pxd | 4 +--- psycopg3/adapt.py | 2 +- psycopg3/adapt.pyx | 38 ++++++++++++++++++++++++++++++++++--- psycopg3/transform.pyx | 29 ++-------------------------- psycopg3/types/composite.py | 2 +- psycopg3/types/text.pxd | 4 ++++ psycopg3/types/text.pyx | 25 +++++++++++++++++++++--- 7 files changed, 66 insertions(+), 38 deletions(-) create mode 100644 psycopg3/types/text.pxd diff --git a/psycopg3/adapt.pxd b/psycopg3/adapt.pxd index 627101ace..4d372a336 100644 --- 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 33cd0af79..2af9a6263 100644 --- 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 aa35c8d91..02cb4e1a5 100644 --- 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 = 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 8e77f787c..2c565fdcb 100644 --- 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 = 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 cf1ed0838..c801e575e 100644 --- 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 index 000000000..81b0425ee --- /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 1fb2fb2b4..b143d4e4e 100644 --- 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 = 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 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) -- 2.47.2