diff --git a/cbor2/encoder.py b/cbor2/encoder.py index 6dad6e60..e150a07f 100644 --- 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 index 701fedc5..a1058d41 100644 --- 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 index 0baf9b6b..d4ac8de4 100644 --- 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 index 56443af6..62e9038f 100644 --- 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;