8000 Add Timestamp support (#382) · eb-emilio/msgpack-python@6414069 · GitHub
[go: up one dir, main page]

Skip to content

Commit 6414069

Browse files
authored
Add Timestamp support (msgpack#382)
1 parent 2c66689 commit 6414069

File tree

9 files changed

+283
-20
lines changed

9 files changed

+283
-20
lines changed

docs/api.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@ API reference
2727

2828
.. autoclass:: ExtType
2929

30+
.. autoclass:: Timestamp
31+
:members:
32+
:special-members: __init__
33+
3034
exceptions
3135
----------
3236

msgpack/__init__.py

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,10 @@
11
# coding: utf-8
22
from ._version import version
33
from .exceptions import *
4+
from .ext import ExtType, Timestamp
45

56
import os
67
import sys
7-
from collections import namedtuple
8-
9-
10-
class ExtType(namedtuple('ExtType', 'code data')):
11-
"""ExtType represents ext type in msgpack."""
12-
def __new__(cls, code, data):
13-
if not isinstance(code, int):
14-
raise TypeError("code must be int")
15-
if not isinstance(data, bytes):
16-
raise TypeError("data must be bytes")
17-
if not 0 <= code <= 127:
18-
raise ValueError("code must be 0~127")
19-
return super(ExtType, cls).__new__(cls, code, data)
208

219

2210
if os.environ.get('MSGPACK_PUREPYTHON') or sys.version_info[0] == 2:

msgpack/_packer.pyx

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@ from cpython cimport *
44
from cpython.bytearray cimport PyByteArray_Check, PyByteArray_CheckExact
55

66
cdef ExtType
7+
cdef Timestamp
78

8-
from . import ExtType
9+
from .ext import ExtType, Timestamp
910

1011

1112
cdef extern from "Python.h":
@@ -36,6 +37,7 @@ cdef extern from "pack.h":
3637
int msgpack_pack_bin(msgpack_packer* pk, size_t l)
3738
int msgpack_pack_raw_body(msgpack_packer* pk, char* body, size_t l)
3839
int msgpack_pack_ext(msgpack_packer* pk, char typecode, size_t l)
40+
int msgpack_pack_timestamp(msgpack_packer* x, long long seconds, unsigned long nanoseconds);
3941
int msgpack_pack_unicode(msgpack_packer* pk, object o, long long limit)
4042

4143
cdef extern from "buff_converter.h":
@@ -135,6 +137,7 @@ cdef class Packer(object):
135137
cdef int _pack(self, object o, int nest_limit=DEFAULT_RECURSE_LIMIT) except -1:
136138
cdef long long llval
137139
cdef unsigned long long ullval
140+
cdef unsigned long ulval
138141
cdef long longval
139142
cdef float fval
140143
cdef double dval
@@ -238,6 +241,10 @@ cdef class Packer(object):
238241
raise ValueError("EXT data is too large")
239242
ret = msgpack_pack_ext(&self.pk, longval, L)
240243
ret = msgpack_pack_raw_body(&self.pk, rawval, L)
244+
elif type(o) is Timestamp:
245+
llval = o.seconds
246+
ulval = o.nanoseconds
247+
ret = msgpack_pack_timestamp(&self.pk, llval, ulval)
241248
elif PyList_CheckExact(o) if strict_types else (PyTuple_Check(o) or PyList_Check(o)):
242249
L = Py_SIZE(o)
243250
if L > ITEM_LIMIT:

msgpack/_unpacker.pyx

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ from .exceptions import (
1919
FormatError,
2020
StackError,
2121
)
22-
from . import ExtType
22+
from .ext import ExtType, Timestamp
2323

2424

2525
cdef extern from "unpack.h":
@@ -31,6 +31,7 @@ cdef extern from "unpack.h":
3131
PyObject* object_hook
3232
PyObject* list_hook
3333
PyObject* ext_hook
34+
PyObject* timestamp_t
3435
char *unicode_errors
3536
Py_ssize_t max_str_len
3637
Py_ssize_t max_bin_len
@@ -98,6 +99,8 @@ cdef inline init_ctx(unpack_context *ctx,
9899
raise TypeError("ext_hook must be a callable.")
99100
ctx.user.ext_hook = <PyObject*>ext_hook
100101

102+
# Add Timestamp type to the user object so it may be used in unpack.h
103+
ctx.user.timestamp_t = <PyObject*>Timestamp
101104
ctx.user.unicode_errors = unicode_errors
102105

103106
def default_read_extended_type(typecode, data):

msgpack/ext.py

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
# coding: utf-8
2+
from collections import namedtuple
3+
import sys
4+
import struct
5+
6+
7+
PY2 = sys.version_info[0] == 2
8+
if not PY2:
9+
long = int
10+
11+
12+
class ExtType(namedtuple('ExtType', 'code data')):
13+
"""ExtType represents ext type in msgpack."""
14+
def __new__(cls, code, data):
15+
if not isinstance(code, int):
16+
raise TypeError("code must be int")
17+
if not isinstance(data, bytes):
18+
raise TypeError("data must be bytes")
19+
if code == -1:
20+
return Timestamp.from_bytes(data)
21+
if not 0 <= code <= 127:
22+
raise ValueError("code must be 0~127")
23+
return super(ExtType, cls).__new__(cls, code, data)
24+
25+
26+
class Timestamp(object):
27+
"""Timestamp represents the Timestamp extension type in msgpack.
28+
29+
When built with Cython, msgpack uses C methods to pack and unpack `Timestamp`. When using pure-Python
30+
msgpack, :func:`to_bytes` and :func:`from_bytes` are used to pack and unpack `Timestamp`.
31+
"""
32+
__slots__ = ["seconds", "nanoseconds"]
33+
34+
def __init__(self, seconds, nanoseconds=0):
35+
"""Initialize a Timestamp object.
36+
37+
:param seconds: Number of seconds since the UNIX epoch (00:00:00 UTC Jan 1 1970, minus leap seconds). May be
38+
negative. If :code:`seconds` includes a fractional part, :code:`nanoseconds` must be 0.
39+
:type seconds: int or float
40+
41+
:param nanoseconds: Number of nanoseconds to add to `seconds` to get fractional time. Maximum is 999_999_999.
42+
Default is 0.
43+
:type nanoseconds: int
44+
45+
Note: Negative times (before the UNIX epoch) are represented as negative seconds + positive ns.
46+
"""
47+
if not isinstance(seconds, (int, long, float)):
48+
raise TypeError("seconds must be numeric")
49+
if not isinstance(nanoseconds, (int, long)):
50+
raise TypeError("nanoseconds must be an integer")
51+
if nanoseconds:
52+
if nanoseconds < 0 or nanoseconds % 1 != 0 or nanoseconds > (1e9 - 1):
53+
raise ValueError("nanoseconds must be a non-negative integer less than 999999999.")
54+
if not isinstance(seconds, (int, long)):
55+
raise ValueError("seconds must be an integer if also providing nanoseconds.")
56+
self.nanoseconds = nanoseconds
57+
else:
58+
# round helps with floating point issues
59+
self.nanoseconds = int(round(seconds % 1 * 1e9, 0))
60+
self.seconds = int(seconds // 1)
61+
62+
def __repr__(self):
63+
"""String representation of Timestamp."""
64+
return "Timestamp(seconds={0}, nanoseconds={1})".format(self.seconds, self.nanoseconds)
65+
66+
def __eq__(self, other):
67+
"""Check for equality with another Timestamp object"""
68+
if type(other) is self.__class__:
69+
return self.seconds == other.seconds and self.nanoseconds == other.nanoseconds
70+
return False
71+
72+
def __ne__(self, other):
73+
"""not-equals method (see :func:`__eq__()`)"""
74+
return not self.__eq__(other)
75+
76+
@staticmethod
77+
def from_bytes(b):
78+
"""Unpack bytes into a `Timestamp` object.
79+
80+
Used for pure-Python msgpack unpacking.
81+
82+
:param b: Payload from msgpack ext message with code -1
83+
:type b: bytes
84+
85+
:returns: Timestamp object unpacked from msgpack ext payload
86+
:rtype: Timestamp
87+
"""
88+
if len(b) == 4:
89+
seconds = struct.unpack("!L", b)[0]
90+
nanoseconds = 0
91+
elif len(b) == 8:
92+
data64 = struct.unpack("!Q", b)[0]
93+
seconds = data64 & 0x00000003ffffffff
94+
nanoseconds = data64 >> 34
95+
elif len(b) == 12:
96+
nanoseconds, seconds = struct.unpack("!Iq", b)
97+
else:
98+
raise ValueError("Timestamp type can only be created from 32, 64, or 96-bit byte objects")
99+
return Timestamp(seconds, nanoseconds)
100+
101+
def to_bytes(self):
102+
"""Pack this Timestamp object into bytes.
103+
104+
Used for pure-Python msgpack packing.
105+
106+
:returns data: Payload for EXT message with code -1 (timestamp type)
107+
:rtype: bytes
108+
"""
109+
if (self.seconds >> 34) == 0: # seconds is non-negative and fits in 34 bits
110+
data64 = self.nanoseconds << 34 | self.seconds
111+
if data64 & 0xffffffff00000000 == 0:
112+
# nanoseconds is zero and seconds < 2**32, so timestamp 32
113+
data = struct.pack("!L", data64)
114+
else:
115+
# timestamp 64
116+
data = struct.pack("!Q", data64)
117+
else:
118+
# timestamp 96
119+
data = struct.pack("!Iq", self.nanoseconds, self.seconds)
120+
return data
121+
122+
def to_float_s(self):
123+
"""Get the timestamp as a floating-point value.
124+
125+
:returns: posix timestamp
126+
:rtype: float
127+
"""
128+
return self.seconds + self.nanoseconds/1e9
129+
130+
def to_unix_ns(self):
131+
"""Get the timestamp as a unixtime in nanoseconds.
132+
133+
:returns: posix timestamp in nanoseconds
134+
:rtype: int
135+
"""
136+
return int(self.seconds * 1e9 + self.nanoseconds)

msgpack/fallback.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def getvalue(self):
6666
StackError,
6767
)
6868

69-
from . import ExtType
69+
from .ext import ExtType, Timestamp
7070

7171

7272
EX_SKIP = 0
@@ -826,9 +826,13 @@ def _pack(self, obj, nest_limit=DEFAULT_RECURSE_LIMIT,
826826
if self._use_float:
827827
return self._buffer.write(struct.pack(">Bf", 0xca, obj))
828828
return self._buffer.write(struct.pack(">Bd", 0xcb, obj))
829-
if check(obj, ExtType):
830-
code = obj.code
831-
data = obj.data
829+
if check(obj, (ExtType, Timestamp)):
830+
if check(obj, Timestamp):
831+
code = -1
832+
data = obj.to_bytes()
833+
else:
834+
code = obj.code
835+
data = obj.data
832836
assert isinstance(code, int)
833837
assert isinstance(data, bytes)
834838
L = len(data)

msgpack/pack_template.h

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -759,6 +759,39 @@ static inline int msgpack_pack_ext(msgpack_packer* x, char typecode, size_t l)
759759

760760
}
761761

762+
/*
763+
* Pack Timestamp extension type. Follows msgpack-c pack_template.h.
764+
*/
765+
static inline int msgpack_pack_timestamp(msgpack_packer* x, int64_t seconds, uint32_t nanoseconds)
766+
{
767+
if ((seconds >> 34) == 0) {
768+
/* seconds is unsigned and fits in 34 bits */
769+
uint64_t data64 = ((uint64_t)nanoseconds << 34) | (uint64_t)seconds;
770+
if ((data64 & 0xffffffff00000000L) == 0) {
771+
/* no nanoseconds and seconds is 32bits or smaller. timestamp32. */
772+
unsigned char buf[4];
773+
uint32_t data32 = (uint32_t)data64;
774+
msgpack_pack_ext(x, -1, 4);
775+
_msgpack_store32(buf, data32);
776+
msgpack_pack_raw_body(x, buf, 4);
777+
} else {
778+
/* timestamp64 */
779+
unsigned char buf[8];
780+
msgpack_pack_ext(x, -1, 8);
781+
_msgpack_store64(buf, data64);
782+
msgpack_pack_raw_body(x, buf, 8);
783+
784+
}
785+
} else {
786+
/* seconds is signed or >34bits */
787+
unsigned char buf[12];
788+
_msgpack_store32(&buf[0], nanoseconds);
789+
_msgpack_store64(&buf[4], seconds);
790+
msgpack_pack_ext(x, -1, 12);
791+
msgpack_pack_raw_body(x, buf, 12);
792+
}
793+
return 0;
794+
}
762795

763796

764797
#undef msgpack_pack_append_buffer

msgpack/unpack.h

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ typedef struct unpack_user {
2727
PyObject *object_hook;
2828
PyObject *list_hook;
2929
PyObject *ext_hook;
30+
PyObject *timestamp_t;
3031
const char *unicode_errors;
3132
Py_ssize_t max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len;
3233
} unpack_user;
@@ -259,6 +260,38 @@ static inline int unpack_callback_bin(unpack_user* u, const char* b, const char*
259260
return 0;
260261
}
261262

263+
typedef struct msgpack_timestamp {
264+
int64_t tv_sec;
265+
uint32_t tv_nsec;
266+
} msgpack_timestamp;
267+
268+
/*
269+
* Unpack ext buffer to a timestamp. Pulled from msgpack-c timestamp.h.
270+
*/
271+
static inline int unpack_timestamp(const char* buf, unsigned int buflen, msgpack_timestamp* ts) {
272+
switch (buflen) {
273+
case 4:
274+
ts->tv_nsec = 0;
275+
{
276+
uint32_t v = _msgpack_load32(uint32_t, buf);
277+
ts->tv_sec = (int64_t)v;
278+
}
279+
return 0;
280+
case 8: {
281+
uint64_t value =_msgpack_load64(uint64_t, buf);
282+
ts->tv_nsec = (uint32_t)(value >> 34);
283+
ts->tv_sec = value & 0x00000003ffffffffLL;
284+
return 0;
285+
}
286+
case 12:
287+
ts->tv_nsec = _msgpack_load32(uint32_t, buf);
288+
ts->tv_sec = _msgpack_load64(int64_t, buf + 4);
289+
return 0;
290+
default:
291+
return -1;
292+
}
293+
}
294+
262295
static inline int unpack_callback_ext(unpack_user* u, const char* base, const char* pos,
263296
unsigned int length, msgpack_unpack_object* o)
264297
{
@@ -273,7 +306,16 @@ static inline int unpack_callback_ext(unpack_user* u, const char* base, const ch
273306
return -1;
274307
}
275308
// length also includes the typecode, so the actual data is length-1
276-
py = PyObject_CallFunction(u->ext_hook, "(iy#)", (int)typecode, pos, (Py_ssize_t)length-1);
309+
if (typecode == -1) {
310+
msgpack_timestamp ts;
311+
if (unpack_timestamp(pos, length-1, &ts) == 0) {
312+
py = PyObject_CallFunction(u->timestamp_t, "(Lk)", ts.tv_sec, ts.tv_nsec);
313+
} else {
314+
py = NULL;
315+
}
316+
} else {
317+
py = PyObject_CallFunction(u->ext_hook, "(iy#)", (int)typecode, pos, (Py_ssize_t)length-1);
318+
}
277319
if (!py)
278320
return -1;
279321
*o = py;

0 commit comments

Comments
 (0)
0