diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/asyncpg/pgproto/codecs/numeric.pyx')
| -rw-r--r-- | .venv/lib/python3.12/site-packages/asyncpg/pgproto/codecs/numeric.pyx | 356 |
1 files changed, 356 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/asyncpg/pgproto/codecs/numeric.pyx b/.venv/lib/python3.12/site-packages/asyncpg/pgproto/codecs/numeric.pyx new file mode 100644 index 00000000..b75d0961 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/asyncpg/pgproto/codecs/numeric.pyx @@ -0,0 +1,356 @@ +# Copyright (C) 2016-present the asyncpg authors and contributors +# <see AUTHORS file> +# +# This module is part of asyncpg and is released under +# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0 + + +from libc.math cimport abs, log10 +from libc.stdio cimport snprintf + +import decimal + +# defined in postgresql/src/backend/utils/adt/numeric.c +DEF DEC_DIGITS = 4 +DEF MAX_DSCALE = 0x3FFF +DEF NUMERIC_POS = 0x0000 +DEF NUMERIC_NEG = 0x4000 +DEF NUMERIC_NAN = 0xC000 +DEF NUMERIC_PINF = 0xD000 +DEF NUMERIC_NINF = 0xF000 + +_Dec = decimal.Decimal + + +cdef numeric_encode_text(CodecContext settings, WriteBuffer buf, obj): + text_encode(settings, buf, str(obj)) + + +cdef numeric_decode_text(CodecContext settings, FRBuffer *buf): + return _Dec(text_decode(settings, buf)) + + +cdef numeric_encode_binary(CodecContext settings, WriteBuffer buf, obj): + cdef: + object dec + object dt + int64_t exponent + int64_t i + int64_t j + tuple pydigits + int64_t num_pydigits + int16_t pgdigit + int64_t num_pgdigits + int16_t dscale + int64_t dweight + int64_t weight + uint16_t sign + int64_t padding_size = 0 + + if isinstance(obj, _Dec): + dec = obj + else: + dec = _Dec(obj) + + dt = dec.as_tuple() + + if dt.exponent == 'n' or dt.exponent == 'N': + # NaN + sign = NUMERIC_NAN + num_pgdigits = 0 + weight = 0 + dscale = 0 + elif dt.exponent == 'F': + # Infinity + if dt.sign: + sign = NUMERIC_NINF + else: + sign = NUMERIC_PINF + num_pgdigits = 0 + weight = 0 + dscale = 0 + else: + exponent = dt.exponent + if exponent < 0 and -exponent > MAX_DSCALE: + raise ValueError( + 'cannot encode Decimal value into numeric: ' + 'exponent is too small') + + if dt.sign: + sign = NUMERIC_NEG + else: + sign = NUMERIC_POS + + pydigits = dt.digits + num_pydigits = len(pydigits) + + dweight = num_pydigits + exponent - 1 + if dweight >= 0: + weight = (dweight + DEC_DIGITS) // DEC_DIGITS - 1 + else: + weight = -((-dweight - 1) // DEC_DIGITS + 1) + + if weight > 2 ** 16 - 1: + raise ValueError( + 'cannot encode Decimal value into numeric: ' + 'exponent is too large') + + padding_size = \ + (weight + 1) * DEC_DIGITS - (dweight + 1) + num_pgdigits = \ + (num_pydigits + padding_size + DEC_DIGITS - 1) // DEC_DIGITS + + if num_pgdigits > 2 ** 16 - 1: + raise ValueError( + 'cannot encode Decimal value into numeric: ' + 'number of digits is too large') + + # Pad decimal digits to provide room for correct Postgres + # digit alignment in the digit computation loop. + pydigits = (0,) * DEC_DIGITS + pydigits + (0,) * DEC_DIGITS + + if exponent < 0: + if -exponent > MAX_DSCALE: + raise ValueError( + 'cannot encode Decimal value into numeric: ' + 'exponent is too small') + dscale = <int16_t>-exponent + else: + dscale = 0 + + buf.write_int32(2 + 2 + 2 + 2 + 2 * <uint16_t>num_pgdigits) + buf.write_int16(<int16_t>num_pgdigits) + buf.write_int16(<int16_t>weight) + buf.write_int16(<int16_t>sign) + buf.write_int16(dscale) + + j = DEC_DIGITS - padding_size + + for i in range(num_pgdigits): + pgdigit = (pydigits[j] * 1000 + pydigits[j + 1] * 100 + + pydigits[j + 2] * 10 + pydigits[j + 3]) + j += DEC_DIGITS + buf.write_int16(pgdigit) + + +# The decoding strategy here is to form a string representation of +# the numeric var, as it is faster than passing an iterable of digits. +# For this reason the below code is pure overhead and is ~25% slower +# than the simple text decoder above. That said, we need the binary +# decoder to support binary COPY with numeric values. +cdef numeric_decode_binary_ex( + CodecContext settings, + FRBuffer *buf, + bint trail_fract_zero, +): + cdef: + uint16_t num_pgdigits = <uint16_t>hton.unpack_int16(frb_read(buf, 2)) + int16_t weight = hton.unpack_int16(frb_read(buf, 2)) + uint16_t sign = <uint16_t>hton.unpack_int16(frb_read(buf, 2)) + uint16_t dscale = <uint16_t>hton.unpack_int16(frb_read(buf, 2)) + int16_t pgdigit0 + ssize_t i + int16_t pgdigit + object pydigits + ssize_t num_pydigits + ssize_t actual_num_pydigits + ssize_t buf_size + int64_t exponent + int64_t abs_exponent + ssize_t exponent_chars + ssize_t front_padding = 0 + ssize_t num_fract_digits + ssize_t trailing_fract_zeros_adj + char smallbuf[_NUMERIC_DECODER_SMALLBUF_SIZE] + char *charbuf + char *bufptr + bint buf_allocated = False + + if sign == NUMERIC_NAN: + # Not-a-number + return _Dec('NaN') + elif sign == NUMERIC_PINF: + # +Infinity + return _Dec('Infinity') + elif sign == NUMERIC_NINF: + # -Infinity + return _Dec('-Infinity') + + if num_pgdigits == 0: + # Zero + return _Dec('0e-' + str(dscale)) + + pgdigit0 = hton.unpack_int16(frb_read(buf, 2)) + if weight >= 0: + if pgdigit0 < 10: + front_padding = 3 + elif pgdigit0 < 100: + front_padding = 2 + elif pgdigit0 < 1000: + front_padding = 1 + + # The number of fractional decimal digits actually encoded in + # base-DEC_DEIGITS digits sent by Postgres. + num_fract_digits = (num_pgdigits - weight - 1) * DEC_DIGITS + + # The trailing zero adjustment necessary to obtain exactly + # dscale number of fractional digits in output. May be negative, + # which indicates that trailing zeros in the last input digit + # should be discarded. + trailing_fract_zeros_adj = dscale - num_fract_digits + + # Maximum possible number of decimal digits in base 10. + # The actual number might be up to 3 digits smaller due to + # leading zeros in first input digit. + num_pydigits = num_pgdigits * DEC_DIGITS + if trailing_fract_zeros_adj > 0: + num_pydigits += trailing_fract_zeros_adj + + # Exponent. + exponent = (weight + 1) * DEC_DIGITS - front_padding + abs_exponent = abs(exponent) + if abs_exponent != 0: + # Number of characters required to render absolute exponent value + # in decimal. + exponent_chars = <ssize_t>log10(<double>abs_exponent) + 1 + else: + exponent_chars = 0 + + # Output buffer size. + buf_size = ( + 1 + # sign + 1 + # leading zero + 1 + # decimal dot + num_pydigits + # digits + 1 + # possible trailing zero padding + 2 + # exponent indicator (E-,E+) + exponent_chars + # exponent + 1 # null terminator char + ) + + if buf_size > _NUMERIC_DECODER_SMALLBUF_SIZE: + charbuf = <char *>cpython.PyMem_Malloc(<size_t>buf_size) + buf_allocated = True + else: + charbuf = smallbuf + + try: + bufptr = charbuf + + if sign == NUMERIC_NEG: + bufptr[0] = b'-' + bufptr += 1 + + bufptr[0] = b'0' + bufptr[1] = b'.' + bufptr += 2 + + if weight >= 0: + bufptr = _unpack_digit_stripping_lzeros(bufptr, pgdigit0) + else: + bufptr = _unpack_digit(bufptr, pgdigit0) + + for i in range(1, num_pgdigits): + pgdigit = hton.unpack_int16(frb_read(buf, 2)) + bufptr = _unpack_digit(bufptr, pgdigit) + + if dscale: + if trailing_fract_zeros_adj > 0: + for i in range(trailing_fract_zeros_adj): + bufptr[i] = <char>b'0' + + # If display scale is _less_ than the number of rendered digits, + # trailing_fract_zeros_adj will be negative and this will strip + # the excess trailing zeros. + bufptr += trailing_fract_zeros_adj + + if trail_fract_zero: + # Check if the number of rendered digits matches the exponent, + # and if so, add another trailing zero, so the result always + # appears with a decimal point. + actual_num_pydigits = bufptr - charbuf - 2 + if sign == NUMERIC_NEG: + actual_num_pydigits -= 1 + + if actual_num_pydigits == abs_exponent: + bufptr[0] = <char>b'0' + bufptr += 1 + + if exponent != 0: + bufptr[0] = b'E' + if exponent < 0: + bufptr[1] = b'-' + else: + bufptr[1] = b'+' + bufptr += 2 + snprintf(bufptr, <size_t>exponent_chars + 1, '%d', + <int>abs_exponent) + bufptr += exponent_chars + + bufptr[0] = 0 + + pydigits = cpythonx.PyUnicode_FromString(charbuf) + + return _Dec(pydigits) + + finally: + if buf_allocated: + cpython.PyMem_Free(charbuf) + + +cdef numeric_decode_binary(CodecContext settings, FRBuffer *buf): + return numeric_decode_binary_ex(settings, buf, False) + + +cdef inline char *_unpack_digit_stripping_lzeros(char *buf, int64_t pgdigit): + cdef: + int64_t d + bint significant + + d = pgdigit // 1000 + significant = (d > 0) + if significant: + pgdigit -= d * 1000 + buf[0] = <char>(d + <int32_t>b'0') + buf += 1 + + d = pgdigit // 100 + significant |= (d > 0) + if significant: + pgdigit -= d * 100 + buf[0] = <char>(d + <int32_t>b'0') + buf += 1 + + d = pgdigit // 10 + significant |= (d > 0) + if significant: + pgdigit -= d * 10 + buf[0] = <char>(d + <int32_t>b'0') + buf += 1 + + buf[0] = <char>(pgdigit + <int32_t>b'0') + buf += 1 + + return buf + + +cdef inline char *_unpack_digit(char *buf, int64_t pgdigit): + cdef: + int64_t d + + d = pgdigit // 1000 + pgdigit -= d * 1000 + buf[0] = <char>(d + <int32_t>b'0') + + d = pgdigit // 100 + pgdigit -= d * 100 + buf[1] = <char>(d + <int32_t>b'0') + + d = pgdigit // 10 + pgdigit -= d * 10 + buf[2] = <char>(d + <int32_t>b'0') + + buf[3] = <char>(pgdigit + <int32_t>b'0') + buf += 4 + + return buf |