From 2e695b7aefc8d5b463ef82f4d180a10df4662df2 Mon Sep 17 00:00:00 2001
From: Georgy Moiseev <georgy.moiseev@corp.mail.ru>
Date: Fri, 2 Sep 2022 12:47:55 +0300
Subject: [PATCH] msgpack: support tzoffset in datetime

Support non-zero tzoffset in datetime extended type. If tzoffset and
tzindex are not specified, return timezone-naive pandas.Timestamp. If
tzoffset is specified, return timezone-aware pandas.Timestamp with
pytz.FixedOffset timezone info.

Part of #204
---
 tarantool/msgpack_ext_types/datetime.py | 29 +++++++++++++++++++++----
 test/suites/test_msgpack_ext_types.py   | 24 ++++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/tarantool/msgpack_ext_types/datetime.py b/tarantool/msgpack_ext_types/datetime.py
index 74bb8be8..480eb7d3 100644
--- a/tarantool/msgpack_ext_types/datetime.py
+++ b/tarantool/msgpack_ext_types/datetime.py
@@ -1,6 +1,7 @@
 import time
 import math
 import pandas
+import pytz
 
 # https://www.tarantool.io/ru/doc/latest/dev_guide/internals/msgpack_extensions/#the-datetime-type
 #
@@ -45,15 +46,30 @@
 NSEC_IN_SEC = 1000000000
 assert isinstance(NSEC_IN_SEC, int)
 
+SEC_IN_MIN = 60
+assert isinstance(SEC_IN_MIN, int)
+
+MIN_IN_DAY = 60 * 24
+assert isinstance(MIN_IN_DAY, int)
+
 def get_int_as_bytes(data, size):
     return data.to_bytes(size, byteorder=BYTEORDER, signed=True)
 
 def encode(obj):
     seconds = obj.value // NSEC_IN_SEC
     nsec = obj.value % NSEC_IN_SEC
+
     tzoffset = 0
     tzindex = 0
 
+    if obj.tz is not None:
+        if obj.tz.zone is not None:
+            raise NotImplementedError
+        else:
+            utc_offset = obj.tz.utcoffset(0)
+            # There is no precision loss since pytz.FixedOffset is in minutes
+            tzoffset = utc_offset.days * MIN_IN_DAY + utc_offset.seconds // SEC_IN_MIN
+
     bytes_buffer = get_int_as_bytes(seconds, SECONDS_SIZE_BYTES)
 
     if (nsec != 0) or (tzoffset != 0) or (tzindex != 0):
@@ -80,9 +96,14 @@ def decode(data):
         tzoffset = 0
         tzindex = 0
 
-    if (tzoffset != 0) or (tzindex != 0):
-        raise NotImplementedError
-
     total_nsec = seconds * NSEC_IN_SEC + nsec
 
-    return pandas.to_datetime(total_nsec, unit='ns')
+    tzinfo = None
+    if (tzindex != 0):
+        raise NotImplementedError
+    elif (tzoffset != 0):
+        tzinfo = pytz.FixedOffset(tzoffset)
+        return pandas.to_datetime(total_nsec, unit='ns').replace(tzinfo=pytz.utc).tz_convert(tzinfo)
+    else:
+        # return timezone-naive pandas.Timestamp
+        return pandas.to_datetime(total_nsec, unit='ns')
diff --git a/test/suites/test_msgpack_ext_types.py b/test/suites/test_msgpack_ext_types.py
index 7eff0651..605ff98d 100644
--- a/test/suites/test_msgpack_ext_types.py
+++ b/test/suites/test_msgpack_ext_types.py
@@ -10,6 +10,7 @@
 import warnings
 import tarantool
 import pandas
+import pytz
 
 from tarantool.msgpack_ext_types.packer import default as packer_default
 from tarantool.msgpack_ext_types.unpacker import ext_hook as unpacker_ext_hook
@@ -544,6 +545,29 @@ def test_UUID_tarantool_encode(self):
             'tarantool': r"datetime.new({year=2022, month=8, day=31, hour=18, min=7, sec=54, " +
                          r"nsec=308543321})",
         },
+        {
+            'python': pandas.Timestamp(year=2022, month=8, day=31, hour=18, minute=7, second=54,
+                                       microsecond=308543, nanosecond=321),
+            'msgpack': (b'\x7a\xa3\x0f\x63\x00\x00\x00\x00\x59\xff\x63\x12\x00\x00\x00\x00'),
+            'tarantool': r"datetime.new({year=2022, month=8, day=31, hour=18, min=7, sec=54, " +
+                         r"nsec=308543321})",
+        },
+        {
+            'python': pandas.Timestamp(year=2022, month=8, day=31, hour=18, minute=7, second=54,
+                                       microsecond=308543, nanosecond=321,
+                                       tzinfo=pytz.FixedOffset(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})",
+        },
+        {
+            'python': pandas.Timestamp(year=2022, month=8, day=31, hour=18, minute=7, second=54,
+                                       microsecond=308543, nanosecond=321,
+                                       tzinfo=pytz.FixedOffset(-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})",
+        },
     ]
 
     def test_datetime_msgpack_decode(self):