Merge branch 'Changaco-patch-1'

Sekenre · Oct 28, 2019 · 59f3946 · 59f3946
2 parents 02bd5a0 + 2d3e4f0
commit 59f3946
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 6 deletions.
diff --git a/cbor2/encoder.py b/cbor2/encoder.py
@@ -92,7 +92,7 @@ def __init__(self, fp, datetime_as_timestamp=False, timezone=None,
         if canonical:
             self._encoders.update(canonical_encoders)
         if date_as_datetime:
-            self._encoders[date] = self.encode_date
+            self._encoders[date] = CBOREncoder.encode_date
 
     def _find_encoder(self, obj_type):
         for type_, enc in list(iteritems(self._encoders)):

diff --git a/source/encoder.c b/source/encoder.c
@@ -111,19 +111,21 @@ CBOREncoder_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
 
 
 // CBOREncoder.__init__(self, fp=None, datetime_as_timestamp=0, timezone=None,
-//                      value_sharing=False, default=None, canonical=False)
+//                      value_sharing=False, default=None, canonical=False,
+//                      date_as_datetime=False)
 int
 CBOREncoder_init(CBOREncoderObject *self, PyObject *args, PyObject *kwargs)
 {
     static char *keywords[] = {
         "fp", "datetime_as_timestamp", "timezone", "value_sharing", "default",
-        "canonical", NULL
+        "canonical", "date_as_datetime", NULL
     };
-    PyObject *tmp, *fp = NULL, *default_handler = NULL, *tz = NULL;
+    PyObject *tmp, *fp = NULL, *default_handler = NULL, *tz = NULL,
+    *date_as_datetime = NULL;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|pOpOB", keywords,
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|pOpOBO", keywords,
                 &fp, &self->timestamp_format, &tz, &self->value_sharing,
-                &default_handler, &self->enc_style))
+                &default_handler, &self->enc_style, &date_as_datetime))
         return -1;
 
     if (_CBOREncoder_set_fp(self, fp, NULL) == -1)
@@ -153,6 +155,15 @@ CBOREncoder_init(CBOREncoderObject *self, PyObject *args, PyObject *kwargs)
                     _CBOR2_str_update, _CBOR2_canonical_encoders, NULL))
             return -1;
     }
+    if (date_as_datetime) {
+        PyObject *encode_date = PyObject_GetAttr((PyObject *) &CBOREncoderType, _CBOR2_str_encode_date);
+        if (!encode_date)
+            return -1;
+        PyObject *datetime_class = PyDateTimeAPI->DateType;
+        if (PyObject_SetItem(self->encoders, datetime_class, encode_date) == -1)
+            return -1;
+        Py_DECREF(encode_date);
+    }
 
     return 0;
 }

diff --git a/source/module.c b/source/module.c
@@ -601,6 +601,7 @@ PyObject *_CBOR2_str_datestr_re = NULL;
 PyObject *_CBOR2_str_Decimal = NULL;
 PyObject *_CBOR2_str_default_encoders = NULL;
 PyObject *_CBOR2_str_denominator = NULL;
+PyObject *_CBOR2_str_encode_date = NULL;
 PyObject *_CBOR2_str_Fraction = NULL;
 PyObject *_CBOR2_str_fromtimestamp = NULL;
 PyObject *_CBOR2_str_FrozenDict = NULL;
@@ -861,6 +862,7 @@ PyInit__cbor2(void)
     INTERN_STRING(Decimal);
     INTERN_STRING(default_encoders);
     INTERN_STRING(denominator);
+    INTERN_STRING(encode_date);
     INTERN_STRING(Fraction);
     INTERN_STRING(fromtimestamp);
     INTERN_STRING(FrozenDict);

diff --git a/source/module.h b/source/module.h
@@ -43,6 +43,7 @@ extern PyObject *_CBOR2_str_datestr_re;
 extern PyObject *_CBOR2_str_Decimal;
 extern PyObject *_CBOR2_str_default_encoders;
 extern PyObject *_CBOR2_str_denominator;
+extern PyObject *_CBOR2_str_encode_date;
 extern PyObject *_CBOR2_str_Fraction;
 extern PyObject *_CBOR2_str_fromtimestamp;
 extern PyObject *_CBOR2_str_FrozenDict;