Support datetime. (#394)

Author: methane

Makefile

@@ -4,7 +4,7 @@ all: cython
 
 .PHONY: black
 black:
-	black msgpack/ test/
+	black msgpack/ test/ setup.py
 
 .PHONY: cython
 cython:

msgpack/_cmsgpack.pyx

@@ -1,4 +1,11 @@
 # coding: utf-8
 #cython: embedsignature=True, c_string_encoding=ascii, language_level=3
+from cpython.datetime cimport import_datetime, datetime_new
+import_datetime()
+
+import datetime
+cdef object utc = datetime.timezone.utc
+cdef object epoch = datetime_new(1970, 1, 1, 0, 0, 0, 0, tz=utc)
+
 include "_packer.pyx"
 include "_unpacker.pyx"

msgpack/_packer.pyx

@@ -2,6 +2,10 @@
 
 from cpython cimport *
 from cpython.bytearray cimport PyByteArray_Check, PyByteArray_CheckExact
+from cpython.datetime cimport (
+    PyDateTime_CheckExact, PyDelta_CheckExact,
+    datetime_tzinfo, timedelta_days, timedelta_seconds, timedelta_microseconds,
+)
 
 cdef ExtType
 cdef Timestamp
@@ -99,8 +103,9 @@ cdef class Packer(object):
     cdef object _berrors
     cdef const char *unicode_errors
     cdef bint strict_types
-    cdef bool use_float
+    cdef bint use_float
     cdef bint autoreset
+    cdef bint datetime
 
     def __cinit__(self):
         cdef int buf_size = 1024*1024
@@ -110,12 +115,13 @@ cdef class Packer(object):
         self.pk.buf_size = buf_size
         self.pk.length = 0
 
-    def __init__(self, *, default=None, unicode_errors=None,
+    def __init__(self, *, default=None,
                  bint use_single_float=False, bint autoreset=True, bint use_bin_type=True,
-                 bint strict_types=False):
+                 bint strict_types=False, bint datetime=False, unicode_errors=None):
         self.use_float = use_single_float
         self.strict_types = strict_types
         self.autoreset = autoreset
+        self.datetime = datetime
         self.pk.use_bin_type = use_bin_type
         if default is not None:
             if not PyCallable_Check(default):
@@ -262,6 +268,13 @@ cdef class Packer(object):
                 if ret == 0:
                     ret = msgpack_pack_raw_body(&self.pk, <char*>view.buf, L)
                 PyBuffer_Release(&view);
+            elif self.datetime and PyDateTime_CheckExact(o) and datetime_tzinfo(o) is not None:
+                delta = o - epoch
+                if not PyDelta_CheckExact(delta):
+                    raise ValueError("failed to calculate delta")
+                llval = timedelta_days(delta) * <long long>(24*60*60) + timedelta_seconds(delta)
+                ulval = timedelta_microseconds(delta) * 1000
+                ret = msgpack_pack_timestamp(&self.pk, llval, ulval)
             elif not default_used and self._default:
                 o = self._default(o)
                 default_used = 1

msgpack/_unpacker.pyx

@@ -1,7 +1,6 @@
 # coding: utf-8
 
 from cpython cimport *
-
 cdef extern from "Python.h":
     ctypedef struct PyObject
     cdef int PyObject_AsReadBuffer(object o, const void** buff, Py_ssize_t* buf_len) except -1
@@ -21,17 +20,22 @@ from .exceptions import (
 )
 from .ext import ExtType, Timestamp
 
+cdef object giga = 1_000_000_000
+
 
 cdef extern from "unpack.h":
     ctypedef struct msgpack_user:
         bint use_list
         bint raw
         bint has_pairs_hook # call object_hook with k-v pairs
         bint strict_map_key
+        int timestamp
         PyObject* object_hook
         PyObject* list_hook
         PyObject* ext_hook
         PyObject* timestamp_t
+        PyObject *giga;
+        PyObject *utc;
         char *unicode_errors
         Py_ssize_t max_str_len
         Py_ssize_t max_bin_len
@@ -57,7 +61,8 @@ cdef extern from "unpack.h":
 cdef inline init_ctx(unpack_context *ctx,
                      object object_hook, object object_pairs_hook,
                      object list_hook, object ext_hook,
-                     bint use_list, bint raw, bint strict_map_key,
+                     bint use_list, bint raw, int timestamp,
+                     bint strict_map_key,
                      const char* unicode_errors,
                      Py_ssize_t max_str_len, Py_ssize_t max_bin_len,
                      Py_ssize_t max_array_len, Py_ssize_t max_map_len,
@@ -99,8 +104,14 @@ cdef inline init_ctx(unpack_context *ctx,
             raise TypeError("ext_hook must be a callable.")
         ctx.user.ext_hook = <PyObject*>ext_hook
 
+    if timestamp < 0 or 3 < timestamp:
+        raise ValueError("timestamp must be 0..3")
+
     # Add Timestamp type to the user object so it may be used in unpack.h
+    ctx.user.timestamp = timestamp
     ctx.user.timestamp_t = <PyObject*>Timestamp
+    ctx.user.giga = <PyObject*>giga
+    ctx.user.utc = <PyObject*>utc
     ctx.user.unicode_errors = unicode_errors
 
 def default_read_extended_type(typecode, data):
@@ -131,7 +142,7 @@ cdef inline int get_data_from_buffer(object obj,
 
 
 def unpackb(object packed, *, object object_hook=None, object list_hook=None,
-            bint use_list=True, bint raw=False, bint strict_map_key=True,
+            bint use_list=True, bint raw=False, int timestamp=0, bint strict_map_key=True,
             unicode_errors=None,
             object_pairs_hook=None, ext_hook=ExtType,
             Py_ssize_t max_str_len=-1,
@@ -179,7 +190,7 @@ def unpackb(object packed, *, object object_hook=None, object list_hook=None,
 
     try:
         init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, ext_hook,
-                 use_list, raw, strict_map_key, cerr,
+                 use_list, raw, timestamp, strict_map_key, cerr,
                  max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len)
         ret = unpack_construct(&ctx, buf, buf_len, &off)
     finally:
@@ -304,7 +315,7 @@ cdef class Unpacker(object):
         self.buf = NULL
 
     def __init__(self, file_like=None, *, Py_ssize_t read_size=0,
-                 bint use_list=True, bint raw=False, bint strict_map_key=True,
+                 bint use_list=True, bint raw=False, int timestamp=0, bint strict_map_key=True,
                  object object_hook=None, object object_pairs_hook=None, object list_hook=None,
                  unicode_errors=None, Py_ssize_t max_buffer_size=100*1024*1024,
                  object ext_hook=ExtType,
@@ -359,7 +370,7 @@ cdef class Unpacker(object):
             cerr = unicode_errors
 
         init_ctx(&self.ctx, object_hook, object_pairs_hook, list_hook,
-                 ext_hook, use_list, raw, strict_map_key, cerr,
+                 ext_hook, use_list, raw, timestamp, strict_map_key, cerr,
                  max_str_len, max_bin_len, max_array_len,
                  max_map_len, max_ext_len)
 

msgpack/ext.py

@@ -1,12 +1,18 @@
 # coding: utf-8
 from collections import namedtuple
+import datetime
 import sys
 import struct
 
 
 PY2 = sys.version_info[0] == 2
+
 if not PY2:
     long = int
+    try:
+        _utc = datetime.timezone.utc
+    except AttributeError:
+        _utc = datetime.timezone(datetime.timedelta(0))
 
 
 class ExtType(namedtuple("ExtType", "code data")):
@@ -131,18 +137,37 @@ def to_bytes(self):
             data = struct.pack("!Iq", self.nanoseconds, self.seconds)
         return data
 
-    def to_float_s(self):
+    def to_float(self):
         """Get the timestamp as a floating-point value.
 
         :returns: posix timestamp
         :rtype: float
         """
         return self.seconds + self.nanoseconds / 1e9
 
+    @staticmethod
+    def from_float(unix_float):
+        seconds = int(unix_float)
+        nanoseconds = int((unix_float % 1) * 1000000000)
+        return Timestamp(seconds, nanoseconds)
+
     def to_unix_ns(self):
         """Get the timestamp as a unixtime in nanoseconds.
 
         :returns: posix timestamp in nanoseconds
         :rtype: int
         """
         return int(self.seconds * 1e9 + self.nanoseconds)
+
+    if not PY2:
+
+        def to_datetime(self):
+            """Get the timestamp as a UTC datetime.
+
+            :rtype: datetime.
+            """
+            return datetime.datetime.fromtimestamp(self.to_float(), _utc)
+
+        @staticmethod
+        def from_datetime(dt):
+            return Timestamp.from_float(dt.timestamp())

msgpack/fallback.py

@@ -1,5 +1,6 @@
 """Fallback pure Python implementation of msgpack"""
 
+from datetime import datetime as _DateTime
 import sys
 import struct
 
@@ -174,6 +175,14 @@ class Unpacker(object):
         If true, unpack msgpack raw to Python bytes.
         Otherwise, unpack to Python str by decoding with UTF-8 encoding (default).
 
+    :param int timestamp:
+        Control how timestamp type is unpacked:
+
+            0 - Tiemstamp
+            1 - float  (Seconds from the EPOCH)
+            2 - int  (Nanoseconds from the EPOCH)
+            3 - datetime.datetime  (UTC).  Python 2 is not supported.
+
     :param bool strict_map_key:
         If true (default), only str or bytes are accepted for map (dict) keys.
 
@@ -248,6 +257,7 @@ def __init__(
         read_size=0,
         use_list=True,
         raw=False,
+        timestamp=0,
         strict_map_key=True,
         object_hook=None,
         object_pairs_hook=None,
@@ -307,6 +317,9 @@ def __init__(
         self._strict_map_key = bool(strict_map_key)
         self._unicode_errors = unicode_errors
         self._use_list = use_list
+        if not (0 <= timestamp <= 3):
+            raise ValueError("timestamp must be 0..3")
+        self._timestamp = timestamp
         self._list_hook = list_hook
         self._object_hook = object_hook
         self._object_pairs_hook = object_pairs_hook
@@ -672,10 +685,21 @@ def _unpack(self, execute=EX_CONSTRUCT):
             else:
                 obj = obj.decode("utf_8", self._unicode_errors)
             return obj
-        if typ == TYPE_EXT:
-            return self._ext_hook(n, bytes(obj))
         if typ == TYPE_BIN:
             return bytes(obj)
+        if typ == TYPE_EXT:
+            if n == -1:  # timestamp
+                ts = Timestamp.from_bytes(bytes(obj))
+                if self._timestamp == 1:
+                    return ts.to_float()
+                elif self._timestamp == 2:
+                    return ts.to_unix_ns()
+                elif self._timestamp == 3:
+                    return ts.to_datetime()
+                else:
+                    return ts
+            else:
+                return self._ext_hook(n, bytes(obj))
         assert typ == TYPE_IMMEDIATE
         return obj
 
@@ -756,6 +780,12 @@ class Packer(object):
         This is useful when trying to implement accurate serialization
         for python types.
 
+    :param bool datetime:
+        If set to true, datetime with tzinfo is packed into Timestamp type.
+        Note that the tzinfo is stripped in the timestamp.
+        You can get UTC datetime with `timestamp=3` option of the Unapcker.
+        (Python 2 is not supported).
+
     :param str unicode_errors:
         The error handler for encoding unicode. (default: 'strict')
         DO NOT USE THIS!!  This option is kept for very specific usage.
@@ -764,18 +794,22 @@ class Packer(object):
     def __init__(
         self,
         default=None,
-        unicode_errors=None,
         use_single_float=False,
         autoreset=True,
         use_bin_type=True,
         strict_types=False,
+        datetime=False,
+        unicode_errors=None,
     ):
         self._strict_types = strict_types
         self._use_float = use_single_float
         self._autoreset = autoreset
         self._use_bin_type = use_bin_type
-        self._unicode_errors = unicode_errors or "strict"
         self._buffer = StringIO()
+        if PY2 and datetime:
+            raise ValueError("datetime is not supported in Python 2")
+        self._datetime = bool(datetime)
+        self._unicode_errors = unicode_errors or "strict"
         if default is not None:
             if not callable(default):
                 raise TypeError("default must be callable")
@@ -891,6 +925,12 @@ def _pack(
                 return self._pack_map_pairs(
                     len(obj), dict_iteritems(obj), nest_limit - 1
                 )
+
+            if self._datetime and check(obj, _DateTime):
+                obj = Timestamp.from_datetime(obj)
+                default_used = 1
+                continue
+
             if not default_used and self._default is not None:
                 obj = self._default(obj)
                 default_used = 1