8000 Support datetime. (#394) · ossdev07/msgpack-python@2186455 · GitHub
[go: up one dir, main page]
Skip to content

Commit 2186455

Browse files
methanemethane
authored
Support datetime. (msgpack#394)
1 parent 5fd6119 commit 2186455

File tree

8 files changed

+222
-24
lines changed

8 files changed

+222
-24
lines changed

‎Makefile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ all: cython
44

55
.PHONY: black
66
black:
7-
black msgpack/ test/
7+
black msgpack/ test/ setup.py
88

99
.PHONY: cython
1010
cython:

‎msgpack/_cmsgpack.pyx

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
11
# coding: utf-8
22
#cython: embedsignature=True, c_string_encoding=ascii, language_level=3
3+
from cpython.datetime cimport import_datetime, datetime_new
4+
import_datetime()
5+
6+
import datetime
7+
cdef object utc = datetime.timezone.utc
8+
cdef object epoch = datetime_new(1970, 1, 1, 0, 0, 0, 0, tz=utc)
9+
310
include "_packer.pyx"
411
include "_unpacker.pyx"

‎msgpack/_packer.pyx

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22

33
from cpython cimport *
44
from cpython.bytearray cimport PyByteArray_Check, PyByteArray_CheckExact
5+
from cpython.datetime cimport (
6+
PyDateTime_CheckExact, PyDelta_CheckExact,
7+
datetime_tzinfo, timedelta_days, timedelta_seconds, timedelta_microseconds,
8+
)
59

610
cdef ExtType
711
cdef Timestamp
@@ -99,8 +103,9 @@ cdef class Packer(object):
99103
cdef object _berrors
100104
cdef const char *unicode_errors
101105
cdef bint strict_types
102-
cdef bool use_float
106+
cdef bint use_float
103107
cdef bint autoreset
108+
cdef bint datetime
104109

105110
def __cinit__(self):
106111
cdef int buf_size = 1024*1024
@@ -110,12 +115,13 @@ cdef class Packer(object):
110115
self.pk.buf_size = buf_size
111116
self.pk.length = 0
112117

113-
def __init__(self, *, default=None, unicode_errors=None,
118+
def __init__(self, *, default=None,
114119
bint use_single_float=False, bint autoreset=True, bint use_bin_type=True,
115-
bint strict_types=False):
120+
bint strict_types=False, bint datetime=False, unicode_errors=None):
116121
self.use_float = use_single_float
117122
self.strict_types = strict_types
118123
self.autoreset = autoreset
124+
self.datetime = datetime
119125
self.pk.use_bin_type = use_bin_type
120126
if default is not None:
121127
if not PyCallable_Check(default):
@@ -262,6 +268,13 @@ cdef class Packer(object):
262268
if ret == 0:
263269
ret = msgpack_pack_raw_body(&self.pk, <char*>view.buf, L)
264270
PyBuffer_Release(&view);
271+
elif self.datetime and PyDateTime_CheckExact(o) and datetime_tzinfo(o) is not None:
272+
delta = o - epoch
273+
if not PyDelta_CheckExact(delta):
274+
raise ValueError("failed to calculate delta")
275+
llval = timedelta_days(delta) * <long long>(24*60*60) + timedelta_seconds(delta)
276+
ulval = timedelta_microseconds(delta) * 1000
277+
ret = msgpack_pack_timestamp(&self.pk, llval, ulval)
265278
elif not default_used and self._default:
266279
o = self._default(o)
267280
default_used = 1

‎msgpack/_unpacker.pyx

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
# coding: utf-8
22

33
from cpython cimport *
4-
54
cdef extern from "Python.h":
65
ctypedef struct PyObject
76
cdef int PyObject_AsReadBuffer(object o, const void** buff, Py_ssize_t* buf_len) except -1
@@ -21,17 +20,22 @@ from .exceptions import (
2120
)
2221
from .ext import ExtType, Timestamp
2322

23+
cdef object giga = 1_000_000_000
24+
2425

2526
cdef extern from "unpack.h":
2627
ctypedef struct msgpack_user:
2728
bint use_list
2829
bint raw
2930
bint has_pairs_hook # call object_hook with k-v pairs
3031
bint strict_map_key
32+
int timestamp
3133
PyObject* object_hook
3234
PyObject* list_hook
3335
PyObject* ext_hook
3436
PyObject* timestamp_t
37+
PyObject *giga;
38+
PyObject *utc;
3539
char *unicode_errors
3640
Py_ssize_t max_str_len
3741
Py_ssize_t max_bin_len
@@ -57,7 +61,8 @@ cdef extern from "unpack.h":
5761
cdef inline init_ctx(unpack_context *ctx,
5862
object object_hook, object object_pairs_hook,
5963
object list_hook, object ext_hook,
60-
bint use_list, bint raw, bint strict_map_key,
64+
bint use_list, bint raw, int timestamp,
65+
bint strict_map_key,
6166
const char* unicode_errors,
6267
Py_ssize_t max_str_len, Py_ssize_t max_bin_len,
6368
Py_ssize_t max_array_len, Py_ssize_t max_map_len,
@@ -99,8 +104,14 @@ cdef inline init_ctx(unpack_context *ctx,
99104
raise TypeError("ext_hook must be a callable.")
100105
ctx.user.ext_hook = <PyObject*>ext_hook
101106

107+
if timestamp < 0 or 3 < timestamp:
108+
raise ValueError("timestamp must be 0..3")
109+
102110
# Add Timestamp type to the user object so it may be used in unpack.h
111+
ctx.user.timestamp = timestamp
103112
ctx.user.timestamp_t = <PyObject*>Timestamp
113+
ctx.user.giga = <PyObject*>giga
114+
ctx.user.utc = <PyObject*>utc
104115
ctx.user.unicode_errors = unicode_errors
105116

106117
def default_read_extended_type(typecode, data):
@@ -131,7 +142,7 @@ cdef inline int get_data_from_buffer(object obj,
131142

132143

133144
def unpackb(object packed, *, object object_hook=None, object list_hook=None,
134-
bint use_list=True, bint raw=False, bint strict_map_key=True,
145+
bint use_list=True, bint raw=False, int timestamp=0, bint strict_map_key=True,
135146
unicode_errors=None,
136147
object_pairs_hook=None, ext_hook=ExtType,
137148
Py_ssize_t max_str_len=-1,
@@ -179,7 +190,7 @@ def unpackb(object packed, *, object object_hook=None, object list_hook=None,
179190

180191
try:
181192
init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, ext_hook,
182-
use_list, raw, strict_map_key, cerr,
193+
use_list, raw, timestamp, strict_map_key, cerr,
183194
max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len)
184195
ret = unpack_construct(&ctx, buf, buf_len, &off)
185196
finally:
@@ -304,7 +315,7 @@ cdef class Unpacker(object):
304315
self.buf = NULL
305316

306317
def __init__(self, file_like=None, *, Py_ssize_t read_size=0,
307-
bint use_list=True, bint raw=False, bint strict_map_key=True,
318+
bint use_list=True, bint raw=False, int timestamp=0, bint strict_map_key=True,
308319
object object_hook=None, object object_pairs_hook=None, object list_hook=None,
309320
unicode_errors=None, Py_ssize_t max_buffer_size=100*1024*1024,
310321
object ext_hook=ExtType,
@@ -359,7 +370,7 @@ cdef class Unpacker(object):
359370
cerr = unicode_errors
360371

361372
init_ctx(&self.ctx, object_hook, object_pairs_hook, list_hook,
362-
ext_hook, use_list, raw, strict_map_key, cerr,
373+
ext_hook, use_list, raw, timestamp, strict_map_key, cerr,
363374
max_str_len, max_bin_len, max_array_len,
364375
max_map_len, max_ext_len)
365376

‎msgpack/ext.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,18 @@
11
# coding: utf-8
22
from collections import namedtuple
3+
import datetime
34
import sys
45
import struct
56

67

78
PY2 = sys.version_info[0] == 2
9+
810
if not PY2:
911
long = int
12+
try:
13+
_utc = datetime.timezone.utc
14+
except AttributeError:
15+
_utc = datetime.timezone(datetime.timedelta(0))
1016

1117

1218
class ExtType(namedtuple("ExtType", "code data")):
@@ -131,18 +137,37 @@ def to_bytes(self):
131137
data = struct.pack("!Iq", self.nanoseconds, self.seconds)
132138
return data
133139

134-
def to_float_s(self):
140+
def to_float(self):
135141
"""Get the timestamp as a floating-point value.
136142
137143
:returns: posix timestamp
138144
:rtype: float
139145
"""
140146
return self.seconds + self.nanoseconds / 1e9
141147

148+
@staticmethod
149+
def from_float(unix_float):
150+
seconds = int(unix_float)
151+
nanoseconds = int((unix_float % 1) * 1000000000)
152+
return Timestamp(seconds, nanoseconds)
153+
142154
def to_unix_ns(self):
143155
"""Get the timestamp as a unixtime in nanoseconds.
144156
145157
:returns: posix timestamp in nanoseconds
146158
:rtype: int
147159
"""
148160
return int(self.seconds * 1e9 + self.nanoseconds)
161+
162+
if not PY2:
163+
164+
def to_datetime(self):
165+
"""Get the timestamp as a UTC datetime.
166+
167+
:rtype: datetime.
168+
"""
169+
return datetime.datetime.fromtimestamp(self.to_float(), _utc)
170+
171+
@staticmethod
172+
def from_datetime(dt):
173+
return Timestamp.from_float(dt.timestamp())

‎msgpack/fallback.py

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""Fallback pure Python implementation of msgpack"""
22

3+
from datetime import datetime as _DateTime
34
import sys
45
import struct
56

@@ -174,6 +175,14 @@ class Unpacker(object):
174175
If true, unpack msgpack raw to Python bytes.
175176
Otherwise, unpack to Python str by decoding with UTF-8 encoding (default).
176177
178+
:param int timestamp:
179+
Control how timestamp type is unpacked:
180+
181+
0 - Tiemstamp
182+
1 - float (Seconds from the EPOCH)
183+
2 - int (Nanoseconds from the EPOCH)
184+
3 - datetime.datetime (UTC). Python 2 is not supported.
185+
177186
:param bool strict_map_key:
178187
If true (default), only str or bytes are accepted for map (dict) keys.
179188
@@ -248,6 +257,7 @@ def __init__(
248257
read_size=0,
249258
use_list=True,
250259
raw=False,
260+
timestamp=0,
251261
strict_map_key=True,
252262
object_hook=None,
253263
object_pairs_hook=None,
@@ -307,6 +317,9 @@ def __init__(
307317
self._strict_map_key = bool(strict_map_key)
308318
self._unicode_errors = unicode_errors
309319
self._use_list = use_list
320+
if not (0 <= timestamp <= 3):
321+
raise ValueError("timestamp must be 0..3")
322+
self._timestamp = timestamp
310323
self._list_hook = list_hook
311324
self._object_hook = object_hook
312325
self._object_pairs_hook = object_pairs_hook
@@ -672,10 +685,21 @@ def _unpack(self, execute=EX_CONSTRUCT):
672685
else:
673686
obj = obj.decode("utf_8", self._unicode_errors)
674687
return obj
675-
if typ == TYPE_EXT:
676-
return self._ext_hook(n, bytes(obj))
677688
if typ == TYPE_BIN:
678689
return bytes(obj)
690+
if typ == TYPE_EXT:
691+
if n == -1: # timestamp
692+
ts = Timestamp.from_bytes(bytes(obj))
693+
if self._timestamp == 1:
694+
return ts.to_float()
695+
elif self._timestamp == 2:
696+
return ts.to_unix_ns()
697+
elif self._timestamp == 3:
698+
return ts.to_datetime()
699+
else:
700+
return ts
701+
else:
702+
return self._ext_hook(n, bytes(obj))
679703
assert typ == TYPE_IMMEDIATE
680704
return obj
681705

@@ -756,6 +780,12 @@ class Packer(object):
756780
This is useful when trying to implement accurate serialization
757781
for python types.
758782
783+
:param bool datetime:
784+
If set to true, datetime with tzinfo is packed into Timestamp type.
785+
Note that the tzinfo is stripped in the timestamp.
786+
You can get UTC datetime with `timestamp=3` option of the Unapcker.
787+
(Python 2 is not supported).
788+
759789
:param str unicode_errors:
760790
The error handler for encoding unicode. (default: 'strict')
761791
DO NOT USE THIS!! This option is kept for very specific usage.
@@ -764,18 +794,22 @@ class Packer(object):
764794
def __init__(
765795
self,
766796
default=None,
767-
unicode_errors=None,
768797
use_single_float=False,
769798
autoreset=True,
770799
use_bin_type=True,
771800
strict_types=False,
801+
datetime=False,
802+
unicode_errors=None,
772803
):
773804
self._strict_types = strict_types
774805
self._use_float = use_single_float
775806
self._autoreset = autoreset
776807
self._use_bin_type = use_bin_type
777-
self._unicode_errors = unicode_errors or "strict"
778808
self._buffer = StringIO()
809+
if PY2 and datetime:
810+
raise ValueError("datetime is not supported in Python 2")
811+
self._datetime = bool(datetime)
812+
self._unicode_errors = unicode_errors or "strict"
779813
if default is not None:
780814
if not callable(default):
781815
raise TypeError("default must be callable")
@@ -891,6 +925,12 @@ def _pack(
891925
return self._pack_map_pairs(
892926
len(obj), dict_iteritems(obj), nest_limit - 1
893927
)
928+
929+
if self._datetime and check(obj, _DateTime):
930+
obj = Timestamp.from_datetime(obj)
931+
default_used = 1
932+
continue
933+
894934
if not default_used and self._default is not None:
895935
obj = self._default(obj)
896936
default_used = 1

0 commit comments

Comments
 (0)
0