msgpack: support tzoffset in datetime

Support non-zero tzoffset in datetime extended type. Use `tzoffset` parameter to set up offset timezone: ``` dt = tarantool.Datetime(year=2022, month=8, day=31, hour=18, minute=7, sec=54, nsec=308543321, tzoffset=180) ``` You may use `tzoffset` property to get timezone offset of a datetime object. Offset timezone is built with pytz.FixedOffset(). pytz module is already a dependency of pandas, but this patch adds it as a requirement just in case something will change in the future. This patch doesn't yet introduce the support of named timezones (tzindex). Part of #204
tarantool · Sep 19, 2022 · e729ac8 · e729ac8
1 parent 623b17d
commit e729ac8
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 14 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -40,6 +40,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
                          nanosecond=(dt.nsec % 1000))
   ```
 
+- Offset in datetime type support (#204).
+
+  Use `tzoffset` parameter to set up offset timezone:
+
+  ```python
+  dt = tarantool.Datetime(year=2022, month=8, day=31,
+                          hour=18, minute=7, sec=54,
+                          nsec=308543321, tzoffset=180)
+  ```
+
+  You may use `tzoffset` property to get timezone offset of a datetime
+  object.
+
 ### Changed
 - Bump msgpack requirement to 1.0.4 (PR #223).
   The only reason of this bump is various vulnerability fixes,

diff --git a/requirements.txt b/requirements.txt
@@ -1,2 +1,3 @@
 msgpack>=1.0.4
 pandas
+pytz
diff --git a/tarantool/msgpack_ext/types/datetime.py b/tarantool/msgpack_ext/types/datetime.py
@@ -1,6 +1,7 @@
 from copy import deepcopy
 
 import pandas
+import pytz
 
 # https://www.tarantool.io/en/doc/latest/dev_guide/internals/msgpack_extensions/#the-datetime-type
 #
@@ -42,6 +43,7 @@
 
 NSEC_IN_SEC = 1000000000
 NSEC_IN_MKSEC = 1000
+SEC_IN_MIN = 60
 
 def get_bytes_as_int(data, cursor, size):
     part = data[cursor:cursor + size]
@@ -50,6 +52,17 @@ def get_bytes_as_int(data, cursor, size):
 def get_int_as_bytes(data, size):
     return data.to_bytes(size, byteorder=BYTEORDER, signed=True)
 
+def compute_offset(timestamp):
+    utc_offset = timestamp.tzinfo.utcoffset(timestamp)
+
+    # `None` offset is a valid utcoffset implementation,
+    # but it seems that pytz timezones never return `None`:
+    # https://github.com/pandas-dev/pandas/issues/15986
+    assert utc_offset is not None
+
+    # There is no precision loss since offset is in minutes
+    return int(utc_offset.total_seconds()) // SEC_IN_MIN
+
 def msgpack_decode(data):
     cursor = 0
     seconds, cursor = get_bytes_as_int(data, cursor, SECONDS_SIZE_BYTES)
@@ -67,16 +80,21 @@ def msgpack_decode(data):
     else:
         raise MsgpackError(f'Unexpected datetime payload length {data_len}')
 
-    if (tzoffset != 0) or (tzindex != 0):
-        raise NotImplementedError
-
     total_nsec = seconds * NSEC_IN_SEC + nsec
+    datetime = pandas.to_datetime(total_nsec, unit='ns')
 
-    return pandas.to_datetime(total_nsec, unit='ns')
+    if tzindex != 0:
+        raise NotImplementedError
+    elif tzoffset != 0:
+        tzinfo = pytz.FixedOffset(tzoffset)
+        return datetime.replace(tzinfo=pytz.UTC).tz_convert(tzinfo)
+    else:
+        return datetime
 
 class Datetime():
     def __init__(self, data=None, *, timestamp=None, year=None, month=None,
-                 day=None, hour=None, minute=None, sec=None, nsec=None):
+                 day=None, hour=None, minute=None, sec=None, nsec=None,
+                 tzoffset=0):
         if data is not None:
             if not isinstance(data, bytes):
                 raise ValueError('data argument (first positional argument) ' +
@@ -99,9 +117,9 @@ def __init__(self, data=None, *, timestamp=None, year=None, month=None,
                     raise ValueError('timestamp must be int if nsec provided')
 
                 total_nsec = timestamp * NSEC_IN_SEC + nsec
-                self._datetime = pandas.to_datetime(total_nsec, unit='ns')
+                datetime = pandas.to_datetime(total_nsec, unit='ns')
             else:
-                self._datetime = pandas.to_datetime(timestamp, unit='s')
+                datetime = pandas.to_datetime(timestamp, unit='s')
         else:
             if nsec is not None:
                 microsecond = nsec // NSEC_IN_MKSEC
@@ -110,10 +128,16 @@ def __init__(self, data=None, *, timestamp=None, year=None, month=None,
                 microsecond = 0
                 nanosecond = 0
 
-            self._datetime = pandas.Timestamp(year=year, month=month, day=day,
-                                              hour=hour, minute=minute, second=sec,
-                                              microsecond=microsecond,
-                                              nanosecond=nanosecond)
+            datetime = pandas.Timestamp(year=year, month=month, day=day,
+                                        hour=hour, minute=minute, second=sec,
+                                        microsecond=microsecond,
+                                        nanosecond=nanosecond)
+
+        if tzoffset != 0:
+            tzinfo = pytz.FixedOffset(tzoffset)
+            datetime = datetime.replace(tzinfo=tzinfo)
+
+        self._datetime = datetime
 
     def __eq__(self, other):
         if isinstance(other, Datetime):
@@ -176,10 +200,16 @@ def nsec(self):
     def timestamp(self):
         return self._datetime.timestamp()
 
+    @property
+    def tzoffset(self):
+        if self._datetime.tzinfo is not None:
+            return compute_offset(self._datetime)
+        return 0
+
     def msgpack_encode(self):
         seconds = self._datetime.value // NSEC_IN_SEC
         nsec = self.nsec
-        tzoffset = 0
+        tzoffset = self.tzoffset
         tzindex = 0
 
         buf = get_int_as_bytes(seconds, SECONDS_SIZE_BYTES)

diff --git a/test/suites/test_datetime.py b/test/suites/test_datetime.py
@@ -53,7 +53,7 @@ def setUp(self):
 
     def test_Datetime_class_API(self):
         dt = tarantool.Datetime(year=2022, month=8, day=31, hour=18, minute=7, sec=54,
-                                nsec=308543321)
+                                nsec=308543321, tzoffset=180)
 
         self.assertEqual(dt.year, 2022)
         self.assertEqual(dt.month, 8)
@@ -62,8 +62,9 @@ def test_Datetime_class_API(self):
         self.assertEqual(dt.minute, 7)
         self.assertEqual(dt.sec, 54)
         self.assertEqual(dt.nsec, 308543321)
-        self.assertEqual(dt.timestamp, 1661969274.308543)
         # Both Tarantool and pandas prone to precision loss for timestamp() floats
+        self.assertEqual(dt.timestamp, 1661958474.308543)
+        self.assertEqual(dt.tzoffset, 180)
 
 
     datetime_class_invalid_init_cases = {
@@ -157,6 +158,30 @@ def test_Datetime_class_invalid_init(self):
             'msgpack': (b'\x7a\xa3\x0f\x63\x00\x00\x00\x00\x59\xff\x63\x12\x00\x00\x00\x00'),
             'tarantool': r"datetime.new({timestamp=1661969274, nsec=308543321})",
         },
+        'datetime_with_positive_offset': {
+            'python': tarantool.Datetime(year=2022, month=8, day=31, hour=18, minute=7, sec=54,
+                                         nsec=308543321, tzoffset=180),
+            'msgpack': (b'\x4a\x79\x0f\x63\x00\x00\x00\x00\x59\xff\x63\x12\xb4\x00\x00\x00'),
+            'tarantool': r"datetime.new({year=2022, month=8, day=31, hour=18, min=7, sec=54, " +
+                         r"nsec=308543321, tzoffset=180})",
+        },
+        'datetime_with_negative_offset': {
+            'python': tarantool.Datetime(year=2022, month=8, day=31, hour=18, minute=7, sec=54,
+                                         nsec=308543321, tzoffset=-60),
+            'msgpack': (b'\x8a\xb1\x0f\x63\x00\x00\x00\x00\x59\xff\x63\x12\xc4\xff\x00\x00'),
+            'tarantool': r"datetime.new({year=2022, month=8, day=31, hour=18, min=7, sec=54, " +
+                         r"nsec=308543321, tzoffset=-60})",
+        },
+        'timestamp_with_positive_offset': {
+            'python': tarantool.Datetime(timestamp=1661969274, tzoffset=180),
+            'msgpack': (b'\x4a\x79\x0f\x63\x00\x00\x00\x00\x00\x00\x00\x00\xb4\x00\x00\x00'),
+            'tarantool': r"datetime.new({timestamp=1661969274, tzoffset=180})",
+        },
+        'timestamp_with_negative_offset': {
+            'python': tarantool.Datetime(timestamp=1661969274, tzoffset=-60),
+            'msgpack': (b'\x8a\xb1\x0f\x63\x00\x00\x00\x00\x00\x00\x00\x00\xc4\xff\x00\x00'),
+            'tarantool': r"datetime.new({timestamp=1661969274, tzoffset=-60})",
+        },
     }
 
     def test_msgpack_decode(self):