diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/asyncpg/protocol/codecs/range.pyx')
| -rw-r--r-- | .venv/lib/python3.12/site-packages/asyncpg/protocol/codecs/range.pyx | 207 |
1 files changed, 207 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/asyncpg/protocol/codecs/range.pyx b/.venv/lib/python3.12/site-packages/asyncpg/protocol/codecs/range.pyx new file mode 100644 index 00000000..1038c18d --- /dev/null +++ b/.venv/lib/python3.12/site-packages/asyncpg/protocol/codecs/range.pyx @@ -0,0 +1,207 @@ +# 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 asyncpg import types as apg_types + +from collections.abc import Sequence as SequenceABC + +# defined in postgresql/src/include/utils/rangetypes.h +DEF RANGE_EMPTY = 0x01 # range is empty +DEF RANGE_LB_INC = 0x02 # lower bound is inclusive +DEF RANGE_UB_INC = 0x04 # upper bound is inclusive +DEF RANGE_LB_INF = 0x08 # lower bound is -infinity +DEF RANGE_UB_INF = 0x10 # upper bound is +infinity + + +cdef enum _RangeArgumentType: + _RANGE_ARGUMENT_INVALID = 0 + _RANGE_ARGUMENT_TUPLE = 1 + _RANGE_ARGUMENT_RANGE = 2 + + +cdef inline bint _range_has_lbound(uint8_t flags): + return not (flags & (RANGE_EMPTY | RANGE_LB_INF)) + + +cdef inline bint _range_has_ubound(uint8_t flags): + return not (flags & (RANGE_EMPTY | RANGE_UB_INF)) + + +cdef inline _RangeArgumentType _range_type(object obj): + if cpython.PyTuple_Check(obj) or cpython.PyList_Check(obj): + return _RANGE_ARGUMENT_TUPLE + elif isinstance(obj, apg_types.Range): + return _RANGE_ARGUMENT_RANGE + else: + return _RANGE_ARGUMENT_INVALID + + +cdef range_encode(ConnectionSettings settings, WriteBuffer buf, + object obj, uint32_t elem_oid, + encode_func_ex encoder, const void *encoder_arg): + cdef: + ssize_t obj_len + uint8_t flags = 0 + object lower = None + object upper = None + WriteBuffer bounds_data = WriteBuffer.new() + _RangeArgumentType arg_type = _range_type(obj) + + if arg_type == _RANGE_ARGUMENT_INVALID: + raise TypeError( + 'list, tuple or Range object expected (got type {})'.format( + type(obj))) + + elif arg_type == _RANGE_ARGUMENT_TUPLE: + obj_len = len(obj) + if obj_len == 2: + lower = obj[0] + upper = obj[1] + + if lower is None: + flags |= RANGE_LB_INF + + if upper is None: + flags |= RANGE_UB_INF + + flags |= RANGE_LB_INC | RANGE_UB_INC + + elif obj_len == 1: + lower = obj[0] + flags |= RANGE_LB_INC | RANGE_UB_INF + + elif obj_len == 0: + flags |= RANGE_EMPTY + + else: + raise ValueError( + 'expected 0, 1 or 2 elements in range (got {})'.format( + obj_len)) + + else: + if obj.isempty: + flags |= RANGE_EMPTY + else: + lower = obj.lower + upper = obj.upper + + if obj.lower_inc: + flags |= RANGE_LB_INC + elif lower is None: + flags |= RANGE_LB_INF + + if obj.upper_inc: + flags |= RANGE_UB_INC + elif upper is None: + flags |= RANGE_UB_INF + + if _range_has_lbound(flags): + encoder(settings, bounds_data, lower, encoder_arg) + + if _range_has_ubound(flags): + encoder(settings, bounds_data, upper, encoder_arg) + + buf.write_int32(1 + bounds_data.len()) + buf.write_byte(<int8_t>flags) + buf.write_buffer(bounds_data) + + +cdef range_decode(ConnectionSettings settings, FRBuffer *buf, + decode_func_ex decoder, const void *decoder_arg): + cdef: + uint8_t flags = <uint8_t>frb_read(buf, 1)[0] + int32_t bound_len + object lower = None + object upper = None + FRBuffer bound_buf + + if _range_has_lbound(flags): + bound_len = hton.unpack_int32(frb_read(buf, 4)) + if bound_len == -1: + lower = None + else: + frb_slice_from(&bound_buf, buf, bound_len) + lower = decoder(settings, &bound_buf, decoder_arg) + + if _range_has_ubound(flags): + bound_len = hton.unpack_int32(frb_read(buf, 4)) + if bound_len == -1: + upper = None + else: + frb_slice_from(&bound_buf, buf, bound_len) + upper = decoder(settings, &bound_buf, decoder_arg) + + return apg_types.Range(lower=lower, upper=upper, + lower_inc=(flags & RANGE_LB_INC) != 0, + upper_inc=(flags & RANGE_UB_INC) != 0, + empty=(flags & RANGE_EMPTY) != 0) + + +cdef multirange_encode(ConnectionSettings settings, WriteBuffer buf, + object obj, uint32_t elem_oid, + encode_func_ex encoder, const void *encoder_arg): + cdef: + WriteBuffer elem_data + ssize_t elem_data_len + ssize_t elem_count + + if not isinstance(obj, SequenceABC): + raise TypeError( + 'expected a sequence (got type {!r})'.format(type(obj).__name__) + ) + + elem_data = WriteBuffer.new() + + for elem in obj: + range_encode(settings, elem_data, elem, elem_oid, encoder, encoder_arg) + + elem_count = len(obj) + if elem_count > INT32_MAX: + raise OverflowError(f'too many elements in multirange value') + + elem_data_len = elem_data.len() + if elem_data_len > INT32_MAX - 4: + raise OverflowError( + f'size of encoded multirange datum exceeds the maximum allowed' + f' {INT32_MAX - 4} bytes') + + # Datum length + buf.write_int32(4 + <int32_t>elem_data_len) + # Number of elements in multirange + buf.write_int32(<int32_t>elem_count) + buf.write_buffer(elem_data) + + +cdef multirange_decode(ConnectionSettings settings, FRBuffer *buf, + decode_func_ex decoder, const void *decoder_arg): + cdef: + int32_t nelems = hton.unpack_int32(frb_read(buf, 4)) + FRBuffer elem_buf + int32_t elem_len + int i + list result + + if nelems == 0: + return [] + + if nelems < 0: + raise exceptions.ProtocolError( + 'unexpected multirange size value: {}'.format(nelems)) + + result = cpython.PyList_New(nelems) + for i in range(nelems): + elem_len = hton.unpack_int32(frb_read(buf, 4)) + if elem_len == -1: + raise exceptions.ProtocolError( + 'unexpected NULL element in multirange value') + else: + frb_slice_from(&elem_buf, buf, elem_len) + elem = range_decode(settings, &elem_buf, decoder, decoder_arg) + cpython.Py_INCREF(elem) + cpython.PyList_SET_ITEM(result, i, elem) + + return result |