Add c implementation of date encoding

agronholm · Jan 2, 2024 · 39489e6 · 39489e6
1 parent e17c4bd
commit 39489e6
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 23 deletions.
diff --git a/source/encoder.c b/source/encoder.c
@@ -108,6 +108,7 @@ CBOREncoder_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
         self->string_references = Py_None;
         self->enc_style = 0;
         self->timestamp_format = false;
+        self->date_as_datetime = false;
         self->value_sharing = false;
         self->shared_handler = NULL;
         self->string_referencing = false;
@@ -139,6 +140,8 @@ CBOREncoder_init(CBOREncoderObject *self, PyObject *args, PyObject *kwargs)
     // Predicate values are returned as ints, but need to be stored as bool or ubyte
     if (timestamp_format == 1)
 	self->timestamp_format = true;
+    if (date_as_datetime == 1)
+    self->date_as_datetime = true;
     if (value_sharing == 1)
 	self->value_sharing = true;
     if (enc_style == 1)
@@ -180,15 +183,6 @@ CBOREncoder_init(CBOREncoderObject *self, PyObject *args, PyObject *kwargs)
                     _CBOR2_str_update, _CBOR2_canonical_encoders, NULL))
             return -1;
     }
-    if (date_as_datetime == 1) {
-        PyObject *encode_date = PyObject_GetAttr((PyObject *) &CBOREncoderType, _CBOR2_str_encode_date);
-        if (!encode_date)
-            return -1;
-        PyObject *datetime_class = (PyObject*)PyDateTimeAPI->DateType;
-        if (PyObject_SetItem(self->encoders, datetime_class, encode_date) == -1)
-            return -1;
-        Py_DECREF(encode_date);
-    }
 
     return 0;
 }
@@ -1043,22 +1037,40 @@ CBOREncoder_encode_datetime(CBOREncoderObject *self, PyObject *value)
 static PyObject *
 CBOREncoder_encode_date(CBOREncoderObject *self, PyObject *value)
 {
-    PyObject *datetime, *ret = NULL;
-
-    if (PyDate_Check(value)) {
-        datetime = PyDateTimeAPI->DateTime_FromDateAndTime(
-                PyDateTime_GET_YEAR(value),
-                PyDateTime_GET_MONTH(value),
-                PyDateTime_GET_DAY(value),
-                0, 0, 0, 0, self->tz,
-                PyDateTimeAPI->DateTimeType);
-        if (datetime) {
-            ret = CBOREncoder_encode_datetime(self, datetime);
-            Py_DECREF(datetime);
-            return ret;
+    // semantic type 1000 or 1004
+
+    PyObject *tmp, *ret = NULL;
+    const char *buf;
+    Py_ssize_t length;
+    Py_INCREF(value);
+    if (value) {
+        if (self->date_as_datetime) {
+            tmp = PyDateTimeAPI->DateTime_FromDateAndTime(
+                    PyDateTime_GET_YEAR(value),
+                    PyDateTime_GET_MONTH(value),
+                    PyDateTime_GET_DAY(value),
+                    0, 0, 0, 0, self->tz,
+                    PyDateTimeAPI->DateTimeType);
+            if (tmp)
+                ret = CBOREncoder_encode_datetime(self, tmp);
         }
+        else if (self->timestamp_format) {
+            tmp = PyObject_CallMethodObjArgs(
+                    value, _CBOR2_str_toordinal, NULL);
+            if (tmp && fp_write(self, "\xD8\x64", 2) == 0) {
+                ret = CBOREncoder_encode_int(self, PyNumber_Subtract(tmp, PyLong_FromLong(719163)));
+            }
+        } else {
+            tmp = PyObject_CallMethodObjArgs(
+                    value, _CBOR2_str_isoformat, NULL);
+            if (tmp && fp_write(self, "\xD9\x03\xEC", 3) == 0) {
+                ret = CBOREncoder_encode_string(self, tmp);
+            }
+        }
+        Py_XDECREF(tmp);
+        Py_DECREF(value);
     }
-    return NULL;
+    return ret;
 }
 
 
@@ -1995,6 +2007,8 @@ encode(CBOREncoderObject *self, PyObject *value)
                 return CBOREncoder_encode_map(self, value);
             else if (PyDateTime_CheckExact(value))
                 return CBOREncoder_encode_datetime(self, value);
+            else if (PyDate_CheckExact(value))
+                return CBOREncoder_encode_date(self, value);
             else if (PyAnySet_CheckExact(value))
                 return CBOREncoder_encode_set(self, value);
             // fall-thru

diff --git a/source/encoder.h b/source/encoder.h
@@ -20,6 +20,7 @@ typedef struct {
     PyObject *shared_handler;
     uint8_t enc_style;  // 0=regular, 1=canonical, 2=custom
     bool timestamp_format;
+    bool date_as_datetime;
     bool value_sharing;
     bool string_referencing;
     bool string_namespacing;

diff --git a/source/module.c b/source/module.c
@@ -642,6 +642,7 @@ PyObject *_CBOR2_str_prefixlen = NULL;
 PyObject *_CBOR2_str_read = NULL;
 PyObject *_CBOR2_str_s = NULL;
 PyObject *_CBOR2_str_timestamp = NULL;
+PyObject *_CBOR2_str_toordinal = NULL;
 PyObject *_CBOR2_str_timezone = NULL;
 PyObject *_CBOR2_str_update = NULL;
 PyObject *_CBOR2_str_utc = NULL;
@@ -973,6 +974,7 @@ PyInit__cbor2(void)
     INTERN_STRING(read);
     INTERN_STRING(s);
     INTERN_STRING(timestamp);
+    INTERN_STRING(toordinal);
     INTERN_STRING(timezone);
     INTERN_STRING(update);
     INTERN_STRING(utc);

diff --git a/source/module.h b/source/module.h
@@ -69,6 +69,7 @@ extern PyObject *_CBOR2_str_prefixlen;
 extern PyObject *_CBOR2_str_read;
 extern PyObject *_CBOR2_str_s;
 extern PyObject *_CBOR2_str_timestamp;
+extern PyObject *_CBOR2_str_toordinal;
 extern PyObject *_CBOR2_str_timezone;
 extern PyObject *_CBOR2_str_update;
 extern PyObject *_CBOR2_str_utc;