pythongh-111306: Improve getter function for datetime fold attribute

Modules/_zoneinfo.c

@@ -166,6 +166,8 @@ free_tzrule(_tzrule *tzrule);
 static PyObject *
 load_timedelta(zoneinfo_state *state, long seconds);
 
+static PyDateTime_DateTime *
+as_datetime(PyObject *dt);
 static int
 get_local_timestamp(PyObject *dt, int64_t *local_ts);
 static _ttinfo *
@@ -2203,7 +2205,8 @@ find_ttinfo(zoneinfo_state *state, PyZoneInfo_ZoneInfo *self, PyObject *dt)
         return NULL;
     }
 
-    unsigned char fold = PyDateTime_DATE_GET_FOLD(dt);
+    PyDateTime_DateTime *dt_dt = as_datetime(dt);
+    unsigned char fold = PyDateTime_DATE_GET_FOLD(dt_dt);
     assert(fold < 2);
     int64_t *local_transitions = self->trans_list_wall[fold];
     size_t num_trans = self->num_transitions;
@@ -2213,7 +2216,7 @@ find_ttinfo(zoneinfo_state *state, PyZoneInfo_ZoneInfo *self, PyObject *dt)
     }
     else if (!num_trans || ts > local_transitions[self->num_transitions - 1]) {
         return find_tzrule_ttinfo(&(self->tzrule_after), ts, fold,
-                                  PyDateTime_GET_YEAR(dt));
+                                  PyDateTime_GET_YEAR(dt_dt));
     }
     else {
         size_t idx = _bisect(ts, local_transitions, self->num_transitions) - 1;
@@ -2243,72 +2246,124 @@ ymd_to_ord(int y, int m, int d)
     return days_before_year + yearday + d;
 }
 
-/* Calculate the number of seconds since 1970-01-01 in local time.
- *
- * This gets a datetime in the same "units" as self->trans_list_wall so that we
- * can easily determine which transitions a datetime falls between. See the
- * comment above ts_to_local for more information.
- * */
-static int
-get_local_timestamp(PyObject *dt, int64_t *local_ts)
+/* Convert a generic PyObject into a PyDateTime_DateTime object. */
+static PyDateTime_DateTime *
+as_datetime(PyObject *dt)
 {
-    assert(local_ts != NULL);
-
+    int year, month, day;
     int hour, minute, second;
-    int ord;
+    unsigned char fold;
     if (PyDateTime_CheckExact(dt)) {
-        int y = PyDateTime_GET_YEAR(dt);
-        int m = PyDateTime_GET_MONTH(dt);
-        int d = PyDateTime_GET_DAY(dt);
+        year = PyDateTime_GET_YEAR(dt);
+        month = PyDateTime_GET_MONTH(dt);
+        day = PyDateTime_GET_DAY(dt);
         hour = PyDateTime_DATE_GET_HOUR(dt);
         minute = PyDateTime_DATE_GET_MINUTE(dt);
         second = PyDateTime_DATE_GET_SECOND(dt);
-
-        ord = ymd_to_ord(y, m, d);
+        fold = PyDateTime_DATE_GET_FOLD(dt);
     }
     else {
-        PyObject *num = PyObject_CallMethod(dt, "toordinal", NULL);
+        PyObject *num = PyObject_CallMethod(dt, "year", NULL);
         if (num == NULL) {
-            return -1;
+            return NULL;
+        }
+        year = PyLong_AsLong(num);
+        Py_DECREF(num);
+        if (year == -1) {
+            return NULL;
         }
 
-        ord = PyLong_AsLong(num);
+        num = PyObject_GetAttrString(dt, "month");
+        if (num == NULL) {
+            return NULL;
+        }
+        month = PyLong_AsLong(num);
         Py_DECREF(num);
-        if (ord == -1 && PyErr_Occurred()) {
-            return -1;
+        if (month == -1) {
+            return NULL;
+        }
+
+        num = PyObject_GetAttrString(dt, "day");
+        if (num == NULL) {
+            return NULL;
+        }
+        day = PyLong_AsLong(num);
+        Py_DECREF(num);
+        if (day == -1) {
+            return NULL;
         }
 
         num = PyObject_GetAttrString(dt, "hour");
         if (num == NULL) {
-            return -1;
+            return NULL;
         }
         hour = PyLong_AsLong(num);
         Py_DECREF(num);
         if (hour == -1) {
-            return -1;
+            return NULL;
         }
 
         num = PyObject_GetAttrString(dt, "minute");
         if (num == NULL) {
-            return -1;
+            return NULL;
         }
         minute = PyLong_AsLong(num);
         Py_DECREF(num);
         if (minute == -1) {
-            return -1;
+            return NULL;
         }
 
         num = PyObject_GetAttrString(dt, "second");
         if (num == NULL) {
-            return -1;
+            return NULL;
         }
         second = PyLong_AsLong(num);
         Py_DECREF(num);
         if (second == -1) {
-            return -1;
+            return NULL;
         }
+
+        num = PyObject_GetAttrString(dt, "fold");
+        if (num == NULL) {
+            return NULL;
+        }
+        int tmp = PyLong_AsLong(num);
+        Py_DECREF(num);
+        if (tmp > 1) {
+            return NULL;
+        }
+        fold = tmp;
+    }
+
+    return (PyDateTime_DateTime *) PyDateTime_FromDateAndTimeAndFold(
+        year, month, day, hour, minute, second, 0, fold
+    );
+}
+
+/* Calculate the number of seconds since 1970-01-01 in local time.
+ *
+ * This gets a datetime in the same "units" as self->trans_list_wall so that we
+ * can easily determine which transitions a datetime falls between. See the
+ * comment above ts_to_local for more information.
+ * */
+static int
+get_local_timestamp(PyObject *dt, int64_t *local_ts)
+{
+    assert(local_ts != NULL);
+
+    PyDateTime_DateTime *dt_dt = as_datetime(dt);
+    if (dt_dt == NULL) {
+        return -1;
     }
 
+    int year = PyDateTime_GET_YEAR(dt_dt);
+    int month = PyDateTime_GET_MONTH(dt_dt);
+    int day = PyDateTime_GET_DAY(dt_dt);
+    int hour = PyDateTime_DATE_GET_HOUR(dt_dt);
+    int minute = PyDateTime_DATE_GET_MINUTE(dt_dt);
+    int second = PyDateTime_DATE_GET_SECOND(dt_dt);
+    int ord = ymd_to_ord(year, month, day);
+
     *local_ts = (int64_t)(ord - EPOCHORDINAL) * 86400L +
                 (int64_t)(hour * 3600L + minute * 60 + second);