msgpack: support datetime interval arithmetic

DifferentialOrange · DifferentialOrange · commit e8b1be00cfd6 · 2022-09-21T15:37:23.000+03:00
Support datetime and interval arithmetic with the same rules as in Tarantool [1]. Valid operations: - `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. - `tarantool.IntervalAdjust.NONE` -- only truncation toward the end of month performed (default mode). ``` >>> 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. ``` >>> 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. ``` >>> 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: "" ``` Tarantool does not yet correctly support subtraction of datetime objects with different timezones [2] and addition of intervals to datetimes with non-fixed offset timezones [3]. tarantool-python implementation support them, but it could be reworked later if core team change another solution. 1. https://github.com/tarantool/tarantool/wiki/Datetime-Internals#interval-arithmetic 2. tarantool/tarantool#7698 3. tarantool/tarantool#7700 Closes #229
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -85,6 +85,52 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   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.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.
+
+  - `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,
diff --git a/tarantool/msgpack_ext/types/datetime.py b/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,78 @@ def __init__(self, data=None, *, timestamp=None, year=None, month=None,
             self._datetime = datetime
             self._tz = ''
 
+    def __add__(self, other):
+        if not isinstance(other, Interval):
+            raise TypeError(f"unsupported operand type(s) for +: '{type(self)}' and '{type(other)}'")
+
+        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 = 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 = other.week,
+                                     days = other.day,
+                                     hours = other.hour,
+                                     minutes = other.minute,
+                                     seconds = other.sec,
+                                     microseconds = other.nsec // NSEC_IN_MKSEC,
+                                     nanoseconds = 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 __sub__(self, other):
+        if not isinstance(other, Datetime):
+            raise TypeError(f"unsupported operand type(s) for -: '{type(self)}' and '{type(other)}'")
+
+        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,
+        )
+
     def __eq__(self, other):
         if isinstance(other, Datetime):
             return self._datetime == other._datetime
diff --git a/tarantool/msgpack_ext/types/interval.py b/tarantool/msgpack_ext/types/interval.py
@@ -102,6 +102,70 @@ def __init__(self, data=None, *, year=0, month=0, week=0,
             self.nsec = nsec
             self.adjust = adjust
 
+    def __add__(self, other):
+        if not isinstance(other, Interval):
+            raise TypeError(f"unsupported operand type(s) for +: '{type(self)}' and '{type(other)}'")
+
+        # Tarantool saves adjust of the first argument
+        #
+        # Tarantool 2.10.1-0-g482d91c66
+        #
+        # tarantool> dt1 = datetime.interval.new{year = 2, adjust='last'}
+        # ---
+        # ...
+        #
+        # tarantool> dt2 = datetime.interval.new{year = 1, adjust='excess'}
+        # ---
+        # ...
+        #
+        # tarantool> (dt1 + dt2).adjust
+        # ---
+        # - 'cdata<enum 112>: 2'
+        # ...
+
+        return Interval(
+            year = self.year + other.year,
+            month = self.month + other.month,
+            day = self.day + other.day,
+            hour = self.hour + other.hour,
+            minute = self.minute + other.minute,
+            sec = self.sec + other.sec,
+            nsec = self.nsec + other.nsec,
+            adjust = self.adjust,
+        )
+
+    def __sub__(self, other):
+        if not isinstance(other, Interval):
+            raise TypeError(f"unsupported operand type(s) for -: '{type(self)}' and '{type(other)}'")
+
+        # Tarantool saves adjust of the first argument
+        #
+        # Tarantool 2.10.1-0-g482d91c66
+        #
+        # tarantool> dt1 = datetime.interval.new{year = 2, adjust='last'}
+        # ---
+        # ...
+        #
+        # tarantool> dt2 = datetime.interval.new{year = 1, adjust='excess'}
+        # ---
+        # ...
+        #
+        # tarantool> (dt1 - dt2).adjust
+        # ---
+        # - 'cdata<enum 112>: 2'
+        # ...
+
+        return Interval(
+            year = self.year - other.year,
+            month = self.month - other.month,
+            day = self.day - other.day,
+            hour = self.hour - other.hour,
+            minute = self.minute - other.minute,
+            sec = self.sec - other.sec,
+            nsec = self.nsec - other.nsec,
+            adjust = self.adjust,
+        )
+
     def __eq__(self, other):
         if not isinstance(other, Interval):
             return False
diff --git a/test/suites/__init__.py b/test/suites/__init__.py
@@ -19,14 +19,16 @@
 from .test_uuid import TestSuite_UUID
 from .test_datetime import TestSuite_Datetime
 from .test_interval import TestSuite_Interval
+from .test_datetime_arithmetic import TestSuite_DatetimeArithmetic
 
 test_cases = (TestSuite_Schema_UnicodeConnection,
               TestSuite_Schema_BinaryConnection,
               TestSuite_Request, TestSuite_Protocol, TestSuite_Reconnect,
               TestSuite_Mesh, TestSuite_Execute, TestSuite_DBAPI,
               TestSuite_Encoding, TestSuite_Pool, TestSuite_Ssl,
               TestSuite_Decimal, TestSuite_UUID, TestSuite_Datetime,
-              TestSuite_Interval)
+              TestSuite_Interval, TestSuite_DatetimeArithmetic)
+
 
 def load_tests(loader, tests, pattern):
     suite = unittest.TestSuite()
diff --git a/test/suites/test_datetime_arithmetic.py b/test/suites/test_datetime_arithmetic.py
