Support datetime interval extended type and arithmetic

CHANGELOG.md

@@ -67,6 +67,73 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
   You may use `tz` property to get timezone name of a datetime object.
 
+- Datetime interval type support and tarantool.Interval type (#229).
+
+  Tarantool datetime interval objects are decoded to `tarantool.Interval`
+  type. `tarantool.Interval` may be encoded to Tarantool interval
+  objects.
+
+  You can create `tarantool.Interval` objects either from msgpack
+  data or by using the same API as in Tarantool:
+
+  ```python
+  di = tarantool.Interval(year=-1, month=2, day=3,
+                          hour=4, minute=-5, sec=6,
+                          nsec=308543321,
+                          adjust=tarantool.IntervalAdjust.NONE)
+  ```
+
+  Its attributes (same as in init API) are exposed, so you can
+  use them if needed.
+
+- Datetime interval arithmetic support (#229).
+
+  Valid operations:
+  - `tarantool.Datetime` + `tarantool.Interval` = `tarantool.Datetime`
+  - `tarantool.Datetime` + `tarantool.Interval` = `tarantool.Datetime`
+  - `tarantool.Datetime` - `tarantool.Datetime` = `tarantool.Interval`
+  - `tarantool.Interval` + `tarantool.Interval` = `tarantool.Interval`
+  - `tarantool.Interval` - `tarantool.Interval` = `tarantool.Interval`
+
+  Since `tarantool.Interval` could contain `month` and `year` fields
+  and such operations could be ambiguous, you can use `adjust` field
+  to tune the logic. The behavior is the same as in Tarantool, see
+  [Interval arithmetic RFC](https://github.com/tarantool/tarantool/wiki/Datetime-Internals#interval-arithmetic).
+
+  - `tarantool.IntervalAdjust.NONE` -- only truncation toward the end of
+    month performed (default mode).
+
+    ```python
+    >>> dt = tarantool.Datetime(year=2022, month=3, day=31)
+    datetime: Timestamp('2022-03-31 00:00:00'), tz: ""
+    >>> di = tarantool.Interval(month=1, adjust=tarantool.IntervalAdjust.NONE)
+    >>> dt + di
+    datetime: Timestamp('2022-04-30 00:00:00'), tz: ""
+    ```
+
+  - `tarantool.IntervalAdjust.EXCESS` -- overflow mode, without any snap
+    or truncation to the end of month, straight addition of days in month,
+    stopping over month boundaries if there is less number of days.
+
+    ```python
+    >>> dt = tarantool.Datetime(year=2022, month=1, day=31)
+    datetime: Timestamp('2022-01-31 00:00:00'), tz: ""
+    >>> di = tarantool.Interval(month=1, adjust=tarantool.IntervalAdjust.EXCESS)
+    >>> dt + di
+    datetime: Timestamp('2022-03-02 00:00:00'), tz: ""
+    ```
+
+  - `tarantool.IntervalAdjust.LAST` -- mode when day snaps to the end of month,
+    if happens.
+
+    ```python
+    >>> dt = tarantool.Datetime(year=2022, month=2, day=28)
+    datetime: Timestamp('2022-02-28 00:00:00'), tz: ""
+    >>> di = tarantool.Interval(month=1, adjust=tarantool.IntervalAdjust.LAST)
+    >>> dt + di
+    datetime: Timestamp('2022-03-31 00:00:00'), tz: ""
+    ```
+
 ### Changed
 - Bump msgpack requirement to 1.0.4 (PR #223).
   The only reason of this bump is various vulnerability fixes,

tarantool/__init__.py

@@ -36,6 +36,11 @@
     Datetime,
 )
 
+from tarantool.msgpack_ext.types.interval import (
+    Adjust as IntervalAdjust,
+    Interval,
+)
+
 __version__ = "0.9.0"
 
 
@@ -95,7 +100,7 @@ def connectmesh(addrs=({'host': 'localhost', 'port': 3301},), user=None,
 
 __all__ = ['connect', 'Connection', 'connectmesh', 'MeshConnection', 'Schema',
            'Error', 'DatabaseError', 'NetworkError', 'NetworkWarning',
-           'SchemaError', 'dbapi', 'Datetime']
+           'SchemaError', 'dbapi', 'Datetime', 'Interval', 'IntervalAdjust']
 
 # ConnectionPool is supported only for Python 3.7 or newer.
 if sys.version_info.major >= 3 and sys.version_info.minor >= 7:

tarantool/msgpack_ext/interval.py

@@ -0,0 +1,9 @@
+from tarantool.msgpack_ext.types.interval import Interval
+
+EXT_ID = 6
+
+def encode(obj):
+    return obj.msgpack_encode()
+
+def decode(data):
+    return Interval(data)

tarantool/msgpack_ext/packer.py

@@ -3,15 +3,18 @@
 from msgpack import ExtType
 
 from tarantool.msgpack_ext.types.datetime import Datetime
+from tarantool.msgpack_ext.types.interval import Interval
 
 import tarantool.msgpack_ext.decimal as ext_decimal
 import tarantool.msgpack_ext.uuid as ext_uuid
 import tarantool.msgpack_ext.datetime as ext_datetime
+import tarantool.msgpack_ext.interval as ext_interval
 
 encoders = [
     {'type': Decimal,  'ext': ext_decimal },
     {'type': UUID,     'ext': ext_uuid    },
     {'type': Datetime, 'ext': ext_datetime},
+    {'type': Interval, 'ext': ext_interval},
 ]
 
 def default(obj):

tarantool/msgpack_ext/types/datetime.py

@@ -6,6 +6,8 @@
 import tarantool.msgpack_ext.types.timezones as tt_timezones
 from tarantool.error import MsgpackError
 
+from tarantool.msgpack_ext.types.interval import Interval, Adjust
+
 # https://www.tarantool.io/en/doc/latest/dev_guide/internals/msgpack_extensions/#the-datetime-type
 #
 # The datetime MessagePack representation looks like this:
@@ -47,6 +49,7 @@
 NSEC_IN_SEC = 1000000000
 NSEC_IN_MKSEC = 1000
 SEC_IN_MIN = 60
+MONTH_IN_YEAR = 12
 
 def get_bytes_as_int(data, cursor, size):
     part = data[cursor:cursor + size]
@@ -168,6 +171,83 @@ def __init__(self, data=None, *, timestamp=None, year=None, month=None,
             self._datetime = datetime
             self._tz = ''
 
+    def _interval_operation(self, other, sign=1):
+        self_dt = self._datetime
+
+        # https://github.com/tarantool/tarantool/wiki/Datetime-Internals#date-adjustions-and-leap-years
+        months = other.year * MONTH_IN_YEAR + other.month
+
+        res = self_dt + pandas.DateOffset(months = sign * months)
+
+        # pandas.DateOffset works exactly like Adjust.NONE
+        if other.adjust == Adjust.EXCESS:
+            if self_dt.day > res.day:
+                res = res + pandas.DateOffset(days = self_dt.day - res.day)
+        elif other.adjust == Adjust.LAST:
+            if self_dt.is_month_end:
+                # day replaces days
+                res = res.replace(day = res.days_in_month)
+
+        res = res + pandas.Timedelta(weeks = sign * other.week,
+                                     days = sign * other.day,
+                                     hours = sign * other.hour,
+                                     minutes = sign * other.minute,
+                                     seconds = sign * other.sec,
+                                     microseconds = sign * (other.nsec // NSEC_IN_MKSEC),
+                                     nanoseconds = sign * (other.nsec % NSEC_IN_MKSEC))
+
+        if res.tzinfo is not None:
+            tzoffset = compute_offset(res)
+        else:
+            tzoffset = 0
+        return Datetime(year=res.year, month=res.month, day=res.day,
+                        hour=res.hour, minute=res.minute, sec=res.second,
+                        nsec=res.nanosecond + res.microsecond * NSEC_IN_MKSEC,
+                        tzoffset=tzoffset, tz=self.tz)
+
+    def __add__(self, other):
+        if not isinstance(other, Interval):
+            raise TypeError(f"unsupported operand type(s) for +: '{type(self)}' and '{type(other)}'")
+
+        return self._interval_operation(other, sign=1)
+
+    def __sub__(self, other):
+        if isinstance(other, Datetime):
+            self_dt = self._datetime
+            other_dt = other._datetime
+
+            # Tarantool datetime subtraction ignores timezone info, but it is a bug:
+            #
+            # Tarantool 2.10.1-0-g482d91c66
+            #
+            # tarantool> datetime.new{tz='MSK'} - datetime.new{tz='UTC'}
+            # ---
+            # - +0 seconds
+            # ...
+            #
+            # Refer to https://github.com/tarantool/tarantool/issues/7698
+            # for possible updates.
+
+            if self_dt.tzinfo != other_dt.tzinfo:
+                other_dt = other_dt.tz_convert(self_dt.tzinfo)
+
+            self_nsec = self_dt.microsecond * NSEC_IN_MKSEC + self_dt.nanosecond
+            other_nsec = other_dt.microsecond * NSEC_IN_MKSEC + other_dt.nanosecond
+
+            return Interval(
+                year = self_dt.year - other_dt.year,
+                month = self_dt.month - other_dt.month,
+                day = self_dt.day - other_dt.day,
+                hour = self_dt.hour - other_dt.hour,
+                minute = self_dt.minute - other_dt.minute,
+                sec = self_dt.second - other_dt.second,
+                nsec = self_nsec - other_nsec,
+            )
+        elif isinstance(other, Interval):
+            return self._interval_operation(other, sign=-1)
+        else:
+            raise TypeError(f"unsupported operand type(s) for -: '{type(self)}' and '{type(other)}'")
+
     def __eq__(self, other):
         if isinstance(other, Datetime):
             return self._datetime == other._datetime