]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
feat: add C numpy dumpers
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Fri, 6 Jan 2023 19:04:53 +0000 (19:04 +0000)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sat, 5 Aug 2023 14:21:30 +0000 (15:21 +0100)
psycopg_c/psycopg_c/_psycopg.pyx
psycopg_c/psycopg_c/types/numeric.pyx
psycopg_c/psycopg_c/types/numpy.pyx [new file with mode: 0644]

index 5bbde2b918359bc4c97cf2dd43b6dd58b2019dcb..f8fb04273274cf25971f59dd4021e7ec93198930 100644 (file)
@@ -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"
index 7c5b05c2617255299941a346086d82dd2fae08b2..fdf9561ac1834b0e0fe44bfb6ea916e19fb83d94 100644 (file)
@@ -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 (file)
index 0000000..9b4a0cb
--- /dev/null
@@ -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)