33import pandas
44import pytz
55
6+ import tarantool .msgpack_ext .types .timezones as tt_timezones
7+ from tarantool .error import MsgpackError
8+
69# https://www.tarantool.io/en/doc/latest/dev_guide/internals/msgpack_extensions/#the-datetime-type
710#
811# The datetime MessagePack representation looks like this:
@@ -63,6 +66,17 @@ def compute_offset(timestamp):
6366 # There is no precision loss since offset is in minutes
6467 return int (utc_offset .total_seconds ()) // SEC_IN_MIN
6568
69+ def get_python_tzinfo (tz , error_class ):
70+ if tz in pytz .all_timezones :
71+ return pytz .timezone (tz )
72+
73+ # Checked with timezones/validate_timezones.py
74+ tt_tzinfo = tt_timezones .timezoneAbbrevInfo [tz ]
75+ if (tt_tzinfo ['category' ] & tt_timezones .TZ_AMBIGUOUS ) != 0 :
76+ raise error_class (f'Failed to create datetime with ambiguous timezone "{ tz } "' )
77+
78+ return pytz .FixedOffset (tt_tzinfo ['offset' ])
79+
6680def msgpack_decode (data ):
6781 cursor = 0
6882 seconds , cursor = get_bytes_as_int (data , cursor , SECONDS_SIZE_BYTES )
@@ -84,23 +98,29 @@ def msgpack_decode(data):
8498 datetime = pandas .to_datetime (total_nsec , unit = 'ns' )
8599
86100 if tzindex != 0 :
87- raise NotImplementedError
101+ if tzindex not in tt_timezones .indexToTimezone :
102+ raise MsgpackError (f'Failed to decode datetime with unknown tzindex "{ tzindex } "' )
103+ tz = tt_timezones .indexToTimezone [tzindex ]
104+ tzinfo = get_python_tzinfo (tz , MsgpackError )
105+ return datetime .replace (tzinfo = pytz .UTC ).tz_convert (tzinfo ), tz
88106 elif tzoffset != 0 :
89107 tzinfo = pytz .FixedOffset (tzoffset )
90- return datetime .replace (tzinfo = pytz .UTC ).tz_convert (tzinfo )
108+ return datetime .replace (tzinfo = pytz .UTC ).tz_convert (tzinfo ), ''
91109 else :
92- return datetime
110+ return datetime , ''
93111
94112class Datetime ():
95113 def __init__ (self , data = None , * , timestamp = None , year = None , month = None ,
96114 day = None , hour = None , minute = None , sec = None , nsec = None ,
97- tzoffset = 0 ):
115+ tzoffset = 0 , tz = '' ):
98116 if data is not None :
99117 if not isinstance (data , bytes ):
100118 raise ValueError ('data argument (first positional argument) ' +
101119 'expected to be a "bytes" instance' )
102120
103- self ._datetime = msgpack_decode (data )
121+ datetime , tz = msgpack_decode (data )
122+ self ._datetime = datetime
123+ self ._tz = tz
104124 return
105125
106126 # The logic is same as in Tarantool, refer to datetime API.
@@ -133,11 +153,20 @@ def __init__(self, data=None, *, timestamp=None, year=None, month=None,
133153 microsecond = microsecond ,
134154 nanosecond = nanosecond )
135155
136- if tzoffset != 0 :
137- tzinfo = pytz . FixedOffset ( tzoffset )
138- datetime = datetime . replace ( tzinfo = tzinfo )
156+ if tz != '' :
157+ if tz not in tt_timezones . timezoneToIndex :
158+ raise ValueError ( f'Unknown Tarantool timezone " { tz } "' )
139159
140- self ._datetime = datetime
160+ tzinfo = get_python_tzinfo (tz , ValueError )
161+ self ._datetime = datetime .replace (tzinfo = tzinfo )
162+ self ._tz = tz
163+ elif tzoffset != 0 :
164+ tzinfo = pytz .FixedOffset (tzoffset )
165+ self ._datetime = datetime .replace (tzinfo = tzinfo )
166+ self ._tz = ''
167+ else :
168+ self ._datetime = datetime
169+ self ._tz = ''
141170
142171 def __eq__ (self , other ):
143172 if isinstance (other , Datetime ):
@@ -151,7 +180,7 @@ def __str__(self):
151180 return self ._datetime .__str__ ()
152181
153182 def __repr__ (self ):
154- return f'datetime: { self ._datetime .__repr__ ()} '
183+ return f'datetime: { self ._datetime .__repr__ ()} , tz: " { self . tz } " '
155184
156185 def __copy__ (self ):
157186 cls = self .__class__
@@ -206,6 +235,10 @@ def tzoffset(self):
206235 return compute_offset (self ._datetime )
207236 return 0
208237
238+ @property
239+ def tz (self ):
240+ return self ._tz
241+
209242 @property
210243 def value (self ):
211244 return self ._datetime .value
@@ -214,7 +247,12 @@ def msgpack_encode(self):
214247 seconds = self .value // NSEC_IN_SEC
215248 nsec = self .nsec
216249 tzoffset = self .tzoffset
217- tzindex = 0
250+
251+ tz = self .tz
252+ if tz != '' :
253+ tzindex = tt_timezones .timezoneToIndex [tz ]
254+ else :
255+ tzindex = 0
218256
219257 buf = get_int_as_bytes (seconds , SECONDS_SIZE_BYTES )
220258
0 commit comments