From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 6 May 2021 17:08:20 +0000 (+0200)
Subject: Add dumper for Python int to numeric
X-Git-Tag: 3.0.dev0~48^2~7
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=2e0b71294d73cc777a2766c4b65c6ae4a7736813;p=thirdparty%2Fpsycopg.git

Add dumper for Python int to numeric
---

diff --git a/psycopg3/psycopg3/types/numeric.py b/psycopg3/psycopg3/types/numeric.py
index fb9b6fea2..e2c1214df 100644
--- a/psycopg3/psycopg3/types/numeric.py
+++ b/psycopg3/psycopg3/types/numeric.py
@@ -5,6 +5,7 @@ Adapers for numeric types.
 # Copyright (C) 2020-2021 The Psycopg Team
 
 import struct
+from math import log
 from typing import Any, Callable, DefaultDict, Dict, Tuple, Union, cast
 from decimal import Decimal, DefaultContext, Context
 
@@ -177,12 +178,33 @@ class Int8BinaryDumper(Int8Dumper):
         return _pack_int8(obj)
 
 
+# Ratio between number of bits required to store a number and number of pg
+# decimal digits required.
+BIT_PER_PGDIGIT = log(2) / log(10_000)
+
+
 class IntNumericBinaryDumper(IntNumericDumper):
 
     format = Format.BINARY
 
-    def dump(self, obj: int) -> bytes:
-        raise NotImplementedError("binary decimal dump not implemented yet")
+    def dump(self, obj: int) -> bytearray:
+        ndigits = int(obj.bit_length() * BIT_PER_PGDIGIT) + 1
+        out = bytearray(b"\x00\x00" * (ndigits + 4))
+        if obj < 0:
+            sign = NUMERIC_NEG
+            obj = -obj
+        else:
+            sign = NUMERIC_POS
+
+        out[:8] = _pack_numeric_head(ndigits, ndigits - 1, sign, 0)
+        i = 8 + (ndigits - 1) * 2
+        while obj:
+            rem = obj % 10_000
+            obj //= 10_000
+            out[i : i + 2] = _pack_uint2(rem)
+            i -= 2
+
+        return out
 
 
 class OidBinaryDumper(OidDumper):
diff --git a/psycopg3_c/psycopg3_c/types/numeric.pyx b/psycopg3_c/psycopg3_c/types/numeric.pyx
index 11786afc3..d835f46e5 100644
--- a/psycopg3_c/psycopg3_c/types/numeric.pyx
+++ b/psycopg3_c/psycopg3_c/types/numeric.pyx
@@ -161,6 +161,14 @@ cdef class Int8BinaryDumper(CDumper):
         return sizeof(int64_t)
 
 
+# Ratio between number of bits required to store a number and number of pg
+# decimal digits required.
+DEF BIT_PER_PGDIGIT = 0.07525749891599529  # log(2) / log(10_000)
+
+DEF NUMERIC_POS = 0x0000
+DEF NUMERIC_NEG = 0x4000
+
+
 @cython.final
 cdef class IntNumericBinaryDumper(CDumper):
 
@@ -170,7 +178,35 @@ cdef class IntNumericBinaryDumper(CDumper):
         self.oid = oids.NUMERIC_OID
 
     cdef Py_ssize_t cdump(self, obj, bytearray rv, Py_ssize_t offset) except -1:
-        raise NotImplementedError("binary decimal dump not implemented yet")
+        # Calculate the number of PG digits required to store the number
+        cdef uint16_t ndigits
+        ndigits = <uint16_t>((<int>obj.bit_length()) * BIT_PER_PGDIGIT) + 1
+
+        cdef uint16_t sign = NUMERIC_POS
+        if obj < 0:
+            sign = NUMERIC_NEG
+            obj = -obj
+
+        cdef Py_ssize_t length = sizeof(uint16_t) * (ndigits + 4)
+        cdef uint16_t *buf
+        buf = <uint16_t *><void *>CDumper.ensure_size(rv, offset, length)
+        buf[0] = endian.htobe16(ndigits)
+        buf[1] = endian.htobe16(ndigits - 1)  # weight
+        buf[2] = endian.htobe16(sign)
+        buf[3] = 0  # dscale
+
+        cdef int i = 4 + ndigits - 1
+        cdef uint16_t rem
+        while obj:
+            rem = obj % 10000
+            obj //= 10000
+            buf[i] = endian.htobe16(rem)
+            i -= 1
+        while i > 3:
+            buf[i] = 0
+            i -= 1
+
+        return length
 
 
 cdef class IntDumper(_NumberDumper):
diff --git a/tests/fix_faker.py b/tests/fix_faker.py
index b5b6f20e5..6770b554b 100644
--- a/tests/fix_faker.py
+++ b/tests/fix_faker.py
@@ -272,7 +272,7 @@ class Faker:
             assert got == want
 
     def make_int(self, spec):
-        return randrange(-(1 << 63), 1 << 63)
+        return randrange(-(1 << 90), 1 << 90)
 
     def make_Int2(self, spec):
         return spec(randrange(-(1 << 15), 1 << 15))
diff --git a/tests/types/test_numeric.py b/tests/types/test_numeric.py
index 3f6f93f82..8fc9ab42f 100644
--- a/tests/types/test_numeric.py
+++ b/tests/types/test_numeric.py
@@ -58,13 +58,16 @@ def test_dump_int(conn, val, expr, fmt_in):
 )
 @pytest.mark.parametrize("fmt_in", [Format.AUTO, Format.TEXT, Format.BINARY])
 def test_dump_int_subtypes(conn, val, expr, fmt_in):
-    if fmt_in in (Format.AUTO, Format.BINARY) and "numeric" in expr:
-        pytest.xfail("binary numeric not implemented")
     cur = conn.cursor()
     cur.execute(f"select pg_typeof({expr}) = pg_typeof(%{fmt_in})", (val,))
     assert cur.fetchone()[0] is True
-    cur.execute(f"select {expr} = %{fmt_in}", (val,))
-    assert cur.fetchone()[0] is True
+    cur.execute(
+        f"select {expr} = %(v){fmt_in}, {expr}::text, %(v){fmt_in}::text",
+        {"v": val},
+    )
+    ok, want, got = cur.fetchone()
+    assert got == want
+    assert ok
 
 
 @pytest.mark.parametrize("fmt_in", [Format.AUTO, Format.TEXT, Format.BINARY])
@@ -89,6 +92,10 @@ def test_dump_enum(conn, fmt_in):
         (-42, b" -42"),
         (int(2 ** 63 - 1), b"9223372036854775807"),
         (int(-(2 ** 63)), b" -9223372036854775808"),
+        (int(2 ** 63), b"9223372036854775808"),
+        (int(-(2 ** 63 + 1)), b" -9223372036854775809"),
+        (int(2 ** 100), b"1267650600228229401496703205376"),
+        (int(-(2 ** 100)), b" -1267650600228229401496703205376"),
     ],
 )
 def test_quote_int(conn, val, expr):