Implement support for range types

edgedb/__init__.py

@@ -21,6 +21,7 @@
 
 from ._version import __version__
 
+from edgedb.datatypes import Range
 from edgedb.datatypes.datatypes import (
     Tuple, NamedTuple, EnumValue, RelativeDuration, ConfigMemory
 )

edgedb/datatypes/__init__.py

@@ -15,3 +15,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
+
+from .range import Range  # noqa

edgedb/datatypes/range.py

@@ -0,0 +1,124 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2022-present MagicStack Inc. and the EdgeDB authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from typing import Any
+
+
+class Range:
+
+    __slots__ = ("_lower", "_upper", "_inc_lower", "_inc_upper", "_empty")
+
+    def __init__(
+        self,
+        lower: Any = None,
+        upper: Any = None,
+        *,
+        inc_lower: bool = True,
+        inc_upper: bool = False,
+        empty: bool = False,
+    ) -> None:
+        self._empty = empty
+
+        if empty:
+            if (
+                lower != upper
+                or lower is not None and inc_upper and inc_lower
+            ):
+                raise ValueError(
+                    "conflicting arguments in range constructor: "
+                    "\"empty\" is `true` while the specified bounds "
+                    "suggest otherwise"
+                )
+
+            self._lower = self._upper = None
+            self._inc_lower = self._inc_upper = False
+        else:
+            self._lower = lower
+            self._upper = upper
+            self._inc_lower = lower is not None and inc_lower
+            self._inc_upper = upper is not None and inc_upper
+
+    @property
+    def lower(self):
+        return self._lower
+
+    @property
+    def inc_lower(self):
+        return self._inc_lower
+
+    @property
+    def upper(self):
+        return self._upper
+
+    @property
+    def inc_upper(self):
+        return self._inc_upper
+
+    def is_empty(self):
+        return self._empty
+
+    def __bool__(self):
+        return not self.is_empty()
+
+    def __eq__(self, other):
+        if not isinstance(other, Range):
+            return NotImplemented
+
+        return (
+            self._lower,
+            self._upper,
+            self._inc_lower,
+            self._inc_upper,
+            self._empty
+        ) == (
+            other._lower,
+            other._upper,
+            other._inc_lower,
+            other._inc_upper,
+            self._empty,
+        )
+
+    def __hash__(self) -> int:
+        return hash((
+            self._lower,
+            self._upper,
+            self._inc_lower,
+            self._inc_upper,
+            self._empty,
+        ))
+
+    def __str__(self) -> str:
+        if self._empty:
+            desc = "empty"
+        else:
+            lb = "(" if not self._inc_lower else "["
+            if self._lower is not None:
+                lb += repr(self._lower)
+
+            if self._upper is not None:
+                ub = repr(self._upper)
+            else:
+                ub = ""
+
+            ub += ")" if self._inc_upper else "]"
+
+            desc = f"{lb}, {ub}"
+
+        return f"<Range {desc}>"
+
+    __repr__ = __str__

edgedb/protocol/codecs/base.pyx

@@ -130,7 +130,10 @@ cdef class BaseRecordCodec(BaseCodec):
     cdef _check_encoder(self):
         if not (self.encoder_flags & RECORD_ENCODER_CHECKED):
             for codec in self.fields_codecs:
-                if not isinstance(codec, (ScalarCodec, ArrayCodec, EnumCodec)):
+                if not isinstance(
+                    codec,
+                    (ScalarCodec, ArrayCodec, EnumCodec, RangeCodec),
+                ):
                     self.encoder_flags |= RECORD_ENCODER_INVALID
                     break
             self.encoder_flags |= RECORD_ENCODER_CHECKED

edgedb/protocol/codecs/codecs.pxd

@@ -23,6 +23,7 @@ include "./tuple.pxd"
 include "./namedtuple.pxd"
 include "./object.pxd"
 include "./array.pxd"
+include "./range.pxd"
 include "./set.pxd"
 include "./enum.pxd"
 

edgedb/protocol/codecs/codecs.pyx

@@ -32,6 +32,7 @@ include "./tuple.pyx"
 include "./namedtuple.pyx"
 include "./object.pyx"
 include "./array.pyx"
+include "./range.pyx"
 include "./set.pyx"
 include "./enum.pyx"
 
@@ -45,6 +46,7 @@ DEF CTYPE_NAMEDTUPLE = 5
 DEF CTYPE_ARRAY = 6
 DEF CTYPE_ENUM = 7
 DEF CTYPE_INPUT_SHAPE = 8
+DEF CTYPE_RANGE = 9
 
 DEF _CODECS_BUILD_CACHE_SIZE = 200
 
@@ -142,6 +144,9 @@ cdef class CodecsRegistry:
                 # First dimension length.
                 frb_read(spec, 4)
 
+            elif t == CTYPE_RANGE:
+                frb_read(spec, 2)
+
             elif t == CTYPE_ENUM:
                 els = <uint16_t>hton.unpack_int16(frb_read(spec, 2))
                 for i in range(els):
@@ -267,6 +272,11 @@ cdef class CodecsRegistry:
             sub_codec = <BaseCodec>codecs_list[pos]
             res = ArrayCodec.new(tid, sub_codec, dim_len)
 
+        elif t == CTYPE_RANGE:
+            pos = <uint16_t>hton.unpack_int16(frb_read(spec, 2))
+            sub_codec = <BaseCodec>codecs_list[pos]
+            res = RangeCodec.new(tid, sub_codec)
+
         else:
             raise NotImplementedError(
                 f'no codec implementation for EdgeDB data class {t}')

edgedb/protocol/codecs/range.pxd

@@ -0,0 +1,27 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2016-present MagicStack Inc. and the EdgeDB authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+@cython.final
+cdef class RangeCodec(BaseCodec):
+
+    cdef:
+        BaseCodec sub_codec
+
+    @staticmethod
+    cdef BaseCodec new(bytes tid, BaseCodec sub_codec)

edgedb/protocol/codecs/range.pyx

@@ -0,0 +1,141 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2016-present MagicStack Inc. and the EdgeDB authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+from edgedb import datatypes
+
+
+cdef uint8_t RANGE_EMPTY = 0x01
+cdef uint8_t RANGE_LB_INC = 0x02
+cdef uint8_t RANGE_UB_INC = 0x04
+cdef uint8_t RANGE_LB_INF = 0x08
+cdef uint8_t RANGE_UB_INF = 0x10
+
+
+@cython.final
+cdef class RangeCodec(BaseCodec):
+
+    def __cinit__(self):
+        self.sub_codec = None
+
+    @staticmethod
+    cdef BaseCodec new(bytes tid, BaseCodec sub_codec):
+        cdef:
+            RangeCodec codec
+
+        codec = RangeCodec.__new__(RangeCodec)
+
+        codec.tid = tid
+        codec.name = 'Range'
+        codec.sub_codec = sub_codec
+
+        return codec
+
+    cdef encode(self, WriteBuffer buf, object obj):
+        cdef:
+            uint8_t flags = 0
+            WriteBuffer sub_data
+            object lower = obj.lower
+            object upper = obj.upper
+            bint inc_lower = obj.inc_lower
+            bint inc_upper = obj.inc_upper
+            bint empty = obj.is_empty()
+
+        if not isinstance(self.sub_codec, ScalarCodec):
+            raise TypeError(
+                'only scalar ranges are supported (got type {!r})'.format(
+                    type(self.sub_codec).__name__
+                )
+            )
+
+        if empty:
+            flags |= RANGE_EMPTY
+        else:
+            if lower is None:
+                flags |= RANGE_LB_INF
+            elif inc_lower:
+                flags |= RANGE_LB_INC
+            if upper is None:
+                flags |= RANGE_UB_INF
+            elif inc_upper:
+                flags |= RANGE_UB_INC
+
+        sub_data = WriteBuffer.new()
+        if lower is not None:
+            try:
+                self.sub_codec.encode(sub_data, lower)
+            except TypeError as e:
+                raise ValueError(
+                    'invalid range lower bound: {}'.format(
+                        e.args[0])) from None
+        if upper is not None:
+            try:
+                self.sub_codec.encode(sub_data, upper)
+            except TypeError as e:
+                raise ValueError(
+                    'invalid range upper bound: {}'.format(
+                        e.args[0])) from None
+
+        buf.write_int32(1 + sub_data.len())
+        buf.write_byte(<int8_t>flags)
+        buf.write_buffer(sub_data)
+
+    cdef decode(self, FRBuffer *buf):
+        cdef:
+            uint8_t flags = <uint8_t>frb_read(buf, 1)[0]
+            bint empty = (flags & RANGE_EMPTY) != 0
+            bint inc_lower = (flags & RANGE_LB_INC) != 0
+            bint inc_upper = (flags & RANGE_UB_INC) != 0
+            bint has_lower = (flags & (RANGE_EMPTY | RANGE_LB_INF)) == 0
+            bint has_upper = (flags & (RANGE_EMPTY | RANGE_UB_INF)) == 0
+            object lower = None
+            object upper = None
+            int32_t sub_len
+            FRBuffer sub_buf
+            BaseCodec sub_codec = self.sub_codec
+
+        if has_lower:
+            sub_len = hton.unpack_int32(frb_read(buf, 4))
+            if sub_len != -1:
+                frb_slice_from(&sub_buf, buf, sub_len)
+                lower = sub_codec.decode(&sub_buf)
+                if frb_get_len(&sub_buf):
+                    raise RuntimeError(
+                        f'unexpected trailing data in buffer after '
+                        f'range bound decoding: {frb_get_len(&sub_buf)}')
+
+        if has_upper:
+            sub_len = hton.unpack_int32(frb_read(buf, 4))
+            if sub_len != -1:
+                frb_slice_from(&sub_buf, buf, sub_len)
+                upper = sub_codec.decode(&sub_buf)
+                if frb_get_len(&sub_buf):
+                    raise RuntimeError(
+                        f'unexpected trailing data in buffer after '
+                        f'range bound decoding: {frb_get_len(&sub_buf)}')
+
+        return datatypes.Range(
+            lower,
+            upper,
+            inc_lower=inc_lower,
+            inc_upper=inc_upper,
+            empty=empty,
+        )
+
+    cdef dump(self, int level = 0):
+        return f'{level * " "}{self.name}\n{self.sub_codec.dump(level + 1)}'