From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 6 Jan 2023 19:04:53 +0000 (+0000)
Subject: feat: add C numpy dumpers
X-Git-Tag: pool-3.2.0~70^2~6
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=e1f9420ca3d1e43cae7a93d6b540f4a3f08d0ab4;p=thirdparty%2Fpsycopg.git

feat: add C numpy dumpers
---

diff --git a/psycopg_c/psycopg_c/_psycopg.pyx b/psycopg_c/psycopg_c/_psycopg.pyx
index 5bbde2b91..f8fb04273 100644
--- a/psycopg_c/psycopg_c/_psycopg.pyx
+++ b/psycopg_c/psycopg_c/_psycopg.pyx
@@ -51,4 +51,5 @@ include "types/array.pyx"
 include "types/datetime.pyx"
 include "types/numeric.pyx"
 include "types/bool.pyx"
+include "types/numpy.pyx"
 include "types/string.pyx"
diff --git a/psycopg_c/psycopg_c/types/numeric.pyx b/psycopg_c/psycopg_c/types/numeric.pyx
index 7c5b05c26..fdf9561ac 100644
--- a/psycopg_c/psycopg_c/types/numeric.pyx
+++ b/psycopg_c/psycopg_c/types/numeric.pyx
@@ -108,13 +108,7 @@ cdef class Int2BinaryDumper(CDumper):
     oid = oids.INT2_OID
 
     cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
-        cdef int16_t *buf = <int16_t *>CDumper.ensure_size(
-            rv, offset, sizeof(int16_t))
-        cdef int16_t val = <int16_t>PyLong_AsLongLong(obj)
-        # swap bytes if needed
-        cdef uint16_t *ptvar = <uint16_t *>(&val)
-        buf[0] = endian.htobe16(ptvar[0])
-        return sizeof(int16_t)
+        return dump_int_to_int2_binary(obj, rv, offset)
 
 
 @cython.final
@@ -124,13 +118,7 @@ cdef class Int4BinaryDumper(CDumper):
     oid = oids.INT4_OID
 
     cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
-        cdef int32_t *buf = <int32_t *>CDumper.ensure_size(
-            rv, offset, sizeof(int32_t))
-        cdef int32_t val = <int32_t>PyLong_AsLongLong(obj)
-        # swap bytes if needed
-        cdef uint32_t *ptvar = <uint32_t *>(&val)
-        buf[0] = endian.htobe32(ptvar[0])
-        return sizeof(int32_t)
+        return dump_int_to_int4_binary(obj, rv, offset)
 
 
 @cython.final
@@ -140,13 +128,7 @@ cdef class Int8BinaryDumper(CDumper):
     oid = oids.INT8_OID
 
     cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
-        cdef int64_t *buf = <int64_t *>CDumper.ensure_size(
-            rv, offset, sizeof(int64_t))
-        cdef int64_t val = PyLong_AsLongLong(obj)
-        # swap bytes if needed
-        cdef uint64_t *ptvar = <uint64_t *>(&val)
-        buf[0] = endian.htobe64(ptvar[0])
-        return sizeof(int64_t)
+        return dump_int_to_int8_binary(obj, rv, offset)
 
 
 cdef extern from *:
@@ -186,6 +168,8 @@ cdef class IntNumericBinaryDumper(CDumper):
 
 cdef class IntDumper(CDumper):
 
+    format = PQ_TEXT
+
     cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
         raise TypeError(
             f"{type(self).__name__} is a dispatcher to other dumpers:"
@@ -713,6 +697,42 @@ cdef Py_ssize_t dump_int_or_sub_to_text(
     return length
 
 
+cdef Py_ssize_t dump_int_to_int2_binary(
+    obj, bytearray rv, Py_ssize_t offset
+) except -1:
+    cdef int16_t *buf = <int16_t *>CDumper.ensure_size(
+        rv, offset, sizeof(int16_t))
+    cdef int16_t val = <int16_t>PyLong_AsLongLong(obj)
+    # swap bytes if needed
+    cdef uint16_t *ptvar = <uint16_t *>(&val)
+    buf[0] = endian.htobe16(ptvar[0])
+    return sizeof(int16_t)
+
+
+cdef Py_ssize_t dump_int_to_int4_binary(
+    obj, bytearray rv, Py_ssize_t offset
+) except -1:
+    cdef int32_t *buf = <int32_t *>CDumper.ensure_size(
+        rv, offset, sizeof(int32_t))
+    cdef int32_t val = <int32_t>PyLong_AsLongLong(obj)
+    # swap bytes if needed
+    cdef uint32_t *ptvar = <uint32_t *>(&val)
+    buf[0] = endian.htobe32(ptvar[0])
+    return sizeof(int32_t)
+
+
+cdef Py_ssize_t dump_int_to_int8_binary(
+    obj, bytearray rv, Py_ssize_t offset
+) except -1:
+    cdef int64_t *buf = <int64_t *>CDumper.ensure_size(
+        rv, offset, sizeof(int64_t))
+    cdef int64_t val = PyLong_AsLongLong(obj)
+    # swap bytes if needed
+    cdef uint64_t *ptvar = <uint64_t *>(&val)
+    buf[0] = endian.htobe64(ptvar[0])
+    return sizeof(int64_t)
+
+
 cdef Py_ssize_t dump_int_to_numeric_binary(obj, bytearray rv, Py_ssize_t offset) except -1:
     # Calculate the number of PG digits required to store the number
     cdef uint16_t ndigits
diff --git a/psycopg_c/psycopg_c/types/numpy.pyx b/psycopg_c/psycopg_c/types/numpy.pyx
new file mode 100644
index 000000000..9b4a0cbc2
--- /dev/null
+++ b/psycopg_c/psycopg_c/types/numpy.pyx
@@ -0,0 +1,71 @@
+"""
+Cython adapters for numpy types.
+"""
+
+# Copyright (C) 2020 The Psycopg Team
+
+cimport cython
+
+
+@cython.final
+cdef class NPInt16Dumper(_IntDumper):
+
+    oid = oids.INT2_OID
+
+
+@cython.final
+cdef class NPInt32Dumper(_IntDumper):
+
+    oid = oids.INT4_OID
+
+
+@cython.final
+cdef class NPInt64Dumper(_IntDumper):
+
+    oid = oids.INT8_OID
+
+
+@cython.final
+cdef class NPNumericDumper(_IntDumper):
+
+    oid = oids.NUMERIC_OID
+
+
+@cython.final
+cdef class NPInt16BinaryDumper(_IntDumper):
+
+    oid = oids.INT2_OID
+    format = PQ_BINARY
+
+    cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
+        return dump_int_to_int2_binary(int(obj), rv, offset)
+
+
+@cython.final
+cdef class NPInt32BinaryDumper(_IntDumper):
+
+    oid = oids.INT4_OID
+    format = PQ_BINARY
+
+    cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
+        return dump_int_to_int4_binary(int(obj), rv, offset)
+
+
+@cython.final
+cdef class NPInt64BinaryDumper(_IntDumper):
+
+    oid = oids.INT8_OID
+    format = PQ_BINARY
+
+    cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
+        return dump_int_to_int8_binary(int(obj), rv, offset)
+
+
+@cython.final
+cdef class NPNumericBinaryDumper(_IntDumper):
+
+    oid = oids.NUMERIC_OID
+    format = PQ_BINARY
+
+    cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
+        return dump_int_to_numeric_binary(int(obj), rv, offset)