1
1
from copy import deepcopy
2
2
3
3
import pandas
4
+ import pytz
4
5
5
6
# https://www.tarantool.io/en/doc/latest/dev_guide/internals/msgpack_extensions/#the-datetime-type
6
7
#
42
43
43
44
NSEC_IN_SEC = 1000000000
44
45
NSEC_IN_MKSEC = 1000
46
+ SEC_IN_MIN = 60
45
47
46
48
def get_bytes_as_int (data , cursor , size ):
47
49
part = data [cursor :cursor + size ]
@@ -50,6 +52,17 @@ def get_bytes_as_int(data, cursor, size):
50
52
def get_int_as_bytes (data , size ):
51
53
return data .to_bytes (size , byteorder = BYTEORDER , signed = True )
52
54
55
+ def compute_offset (timestamp ):
56
+ utc_offset = timestamp .tzinfo .utcoffset (timestamp )
57
+
58
+ # `None` offset is a valid utcoffset implementation,
59
+ # but it seems that pytz timezones never return `None`:
60
+ # https://github.com/pandas-dev/pandas/issues/15986
61
+ assert utc_offset is not None
62
+
63
+ # There is no precision loss since offset is in minutes
64
+ return int (utc_offset .total_seconds ()) // SEC_IN_MIN
65
+
53
66
def msgpack_decode (data ):
54
67
cursor = 0
55
68
seconds , cursor = get_bytes_as_int (data , cursor , SECONDS_SIZE_BYTES )
@@ -67,16 +80,21 @@ def msgpack_decode(data):
67
80
else :
68
81
raise MsgpackError (f'Unexpected datetime payload length { data_len } ' )
69
82
70
- if (tzoffset != 0 ) or (tzindex != 0 ):
71
- raise NotImplementedError
72
-
73
83
total_nsec = seconds * NSEC_IN_SEC + nsec
84
+ datetime = pandas .to_datetime (total_nsec , unit = 'ns' )
74
85
75
- return pandas .to_datetime (total_nsec , unit = 'ns' )
86
+ if tzindex != 0 :
87
+ raise NotImplementedError
88
+ elif tzoffset != 0 :
89
+ tzinfo = pytz .FixedOffset (tzoffset )
90
+ return datetime .replace (tzinfo = pytz .UTC ).tz_convert (tzinfo )
91
+ else :
92
+ return datetime
76
93
77
94
class Datetime ():
78
95
def __init__ (self , data = None , * , timestamp = None , year = None , month = None ,
79
- day = None , hour = None , minute = None , sec = None , nsec = None ):
96
+ day = None , hour = None , minute = None , sec = None , nsec = None ,
97
+ tzoffset = 0 ):
80
98
if data is not None :
81
99
if not isinstance (data , bytes ):
82
100
raise ValueError ('data argument (first positional argument) ' +
@@ -99,9 +117,9 @@ def __init__(self, data=None, *, timestamp=None, year=None, month=None,
99
117
raise ValueError ('timestamp must be int if nsec provided' )
100
118
101
119
total_nsec = timestamp * NSEC_IN_SEC + nsec
102
- self . _datetime = pandas .to_datetime (total_nsec , unit = 'ns' )
120
+ datetime = pandas .to_datetime (total_nsec , unit = 'ns' )
103
121
else :
104
- self . _datetime = pandas .to_datetime (timestamp , unit = 's' )
122
+ datetime = pandas .to_datetime (timestamp , unit = 's' )
105
123
else :
106
124
if nsec is not None :
107
125
microsecond = nsec // NSEC_IN_MKSEC
@@ -110,10 +128,16 @@ def __init__(self, data=None, *, timestamp=None, year=None, month=None,
110
128
microsecond = 0
111
129
nanosecond = 0
112
130
113
- self ._datetime = pandas .Timestamp (year = year , month = month , day = day ,
114
- hour = hour , minute = minute , second = sec ,
115
- microsecond = microsecond ,
116
- nanosecond = nanosecond )
131
+ datetime = pandas .Timestamp (year = year , month = month , day = day ,
132
+ hour = hour , minute = minute , second = sec ,
133
+ microsecond = microsecond ,
134
+ nanosecond = nanosecond )
135
+
136
+ if tzoffset != 0 :
137
+ tzinfo = pytz .FixedOffset (tzoffset )
138
+ datetime = datetime .replace (tzinfo = tzinfo )
139
+
140
+ self ._datetime = datetime
117
141
118
142
def __eq__ (self , other ):
119
143
if isinstance (other , Datetime ):
@@ -176,14 +200,20 @@ def nsec(self):
176
200
def timestamp (self ):
177
201
return self ._datetime .timestamp ()
178
202
203
+ @property
204
+ def tzoffset (self ):
205
+ if self ._datetime .tzinfo is not None :
206
+ return compute_offset (self ._datetime )
207
+ return 0
208
+
179
209
@property
180
210
def value (self ):
181
211
return self ._datetime .value
182
212
183
213
def msgpack_encode (self ):
184
214
seconds = self .value // NSEC_IN_SEC
185
215
nsec = self .nsec
186
- tzoffset = 0
216
+ tzoffset = self . tzoffset
187
217
tzindex = 0
188
218
189
219
buf = get_int_as_bytes (seconds , SECONDS_SIZE_BYTES )
0 commit comments