Support casting of datetime with timezone information

Before these changes, nanobind would silently accept datetime objects
with timezone information and treat them as local time.

Author: lbonn

docs/api_extra.rst

@@ -1164,10 +1164,11 @@ to a Python :py:class:`~datetime.datetime` object; other clocks may convert to
 
 The first clock defined by the standard is ``std::chrono::system_clock``.
 This clock measures the current date and time, much like the Python
-:py:func:`time.time` function. It can change abruptly due to
-administrative actions, daylight savings time transitions, or
-synchronization with an external time server. That makes this clock a
-poor choice for timing purposes, but a good choice for wall-clock time.
+:py:func:`time.time` function. While it tracks elapsed time since 1970, in
+the UTC timezone, It can change abruptly due to administrative actions,
+leap seconds changes or synchronization with an external time server.
+That makes this clock a poor choice for timing purposes,
+but a good choice for wall-clock time.
 
 The second clock defined by the standard is ``std::chrono::steady_clock``.
 This clock ticks at a steady rate and is never adjusted, like
@@ -1204,6 +1205,8 @@ converting to Python.
     :py:class:`~datetime.datetime` instance.  The result describes a time in the
     local timezone, but does not have any timezone information
     attached to it (it is a naive datetime object).
+    Note that some time conversion can occur if your system is not set to use the
+    UTC timezone as the C++ and Python types follow different conventions.
 
 - ``std::chrono::duration`` → :py:class:`datetime.timedelta`
     A duration will be converted to a Python :py:class:`~datetime.timedelta`.
@@ -1219,9 +1222,11 @@ converting to Python.
 - :py:class:`datetime.datetime` or :py:class:`datetime.date` or :py:class:`datetime.time` → ``std::chrono::system_clock::time_point``
     A Python date, time, or datetime object can be converted into a
     system clock timepoint.  A :py:class:`~datetime.time` with no date
-    information is treated as that time on January 1, 1970. A
-    :py:class:`~datetime.date` with no time information is treated as midnight
-    on that date. **Any timezone information is ignored.**
+    information is treated as that time on January 1, 1970 (timezone information
+    is ignored). A :py:class:`~datetime.date` with no time information is treated
+    as midnight on that date. A :py:class:`~datetime.datetime` is treated as a
+    local time if no timezone information is provided. Otherwise, it is treated as
+    a time in its associated tzinfo.
 
 - :py:class:`datetime.timedelta` → ``std::chrono::duration``
     A Python time delta object can be converted into a duration

include/nanobind/stl/chrono.h

@@ -153,8 +153,13 @@ class type_caster<std::chrono::time_point<std::chrono::system_clock, Duration>>
         ch::microseconds msecs;
         int yy, mon, dd, hh, min, ss, uu;
         try {
-            if (!unpack_datetime(src.ptr(), &yy, &mon, &dd,
-                                 &hh, &min, &ss, &uu)) {
+            PyObject* shifted_time_ptr{};
+            if (!shift_to_timezone(src.ptr(), &shifted_time_ptr)) {
+                return false;
+            }
+
+            if (!unpack_datetime(bool(shifted_time_ptr) ? steal(shifted_time_ptr).ptr() : src.ptr(),
+                                 &yy, &mon, &dd, &hh, &min, &ss, &uu)) {
                 return false;
             }
         } catch (python_error& e) {

include/nanobind/stl/detail/chrono.h

@@ -144,6 +144,26 @@ NB_NOINLINE inline bool unpack_timedelta(PyObject *o, int *days,
     return false;
 }
 
+NB_NOINLINE inline bool shift_to_timezone(PyObject *o, PyObject** shifted_o, PyObject* tz = Py_None) {
+    *shifted_o = nullptr;
+
+    datetime_types.ensure_ready();
+    if (!PyType_IsSubtype(Py_TYPE(o), (PyTypeObject*) datetime_types.datetime.ptr()) ||
+        PyObject_GetAttrString(o, "tzinfo") == Py_None) {
+        return true;
+    }
+
+    try {
+        handle shifted_datetime = borrow(o).attr("astimezone")(borrow(tz)).release();
+        *shifted_o = shifted_datetime.ptr();
+    } catch (python_error& e) {
+        e.restore();
+        return false;
+    }
+
+    return true;
+}
+
 NB_NOINLINE inline bool unpack_datetime(PyObject *o,
                                         int *year, int *month, int *day,
                                         int *hour, int *minute, int *second,
@@ -229,6 +249,41 @@ NB_NOINLINE inline bool unpack_timedelta(PyObject *o, int *days,
     return false;
 }
 
+NB_NOINLINE inline bool shift_to_timezone(PyObject *o, PyObject** shifted_o, PyObject* tz = Py_None) {
+    if (!PyDateTimeAPI) {
+        PyDateTime_IMPORT;
+        if (!PyDateTimeAPI)
+            raise_python_error();
+    }
+
+    *shifted_o = nullptr;
+
+    if (!PyDateTime_Check(o)) {
+        return true;
+    }
+
+#if PY_VERSION_HEX < 0x03090000
+    object tzinfo = borrow(o).attr("tzinfo");
+    if (tzinfo.is_none()) {
+        return true;
+    }
+#else
+    if (PyDateTime_DATE_GET_TZINFO(o) == Py_None) {
+        return true;
+    }
+#endif
+
+    try {
+        handle shifted_datetime = borrow(o).attr("astimezone")(borrow(tz)).release();
+        *shifted_o = shifted_datetime.ptr();
+    } catch (python_error& e) {
+        e.restore();
+        return false;
+    }
+
+    return true;
+}
+
 NB_NOINLINE inline bool unpack_datetime(PyObject *o,
                                         int *year, int *month, int *day,
                                         int *hour, int *minute, int *second,

tests/test_chrono.py

@@ -8,6 +8,13 @@
 
 import pytest
 
+# tzset is unix-only but necessary after changing timezones
+try:
+    from time import tzset
+except ImportError:
+    def tzset():
+        pass
+
 
 def test_chrono_system_clock():
     # Get the time from both c++ and datetime
@@ -73,6 +80,12 @@ def test_chrono_system_clock_roundtrip_date():
     assert time2.microsecond == 0
 
 
+SKIP_TZ_NO_ZONEINFO = pytest.mark.skipif(
+    "sys.version_info < (3, 9)",
+    reason="requires python3.9 or higher"
+)
+
+
 SKIP_TZ_ENV_ON_WIN = pytest.mark.skipif(
     "sys.platform == 'win32'",
     reason="TZ environment variable only supported on POSIX"
@@ -104,6 +117,7 @@ def test_chrono_system_clock_roundtrip_date():
 def test_chrono_system_clock_roundtrip_time(time1, tz, monkeypatch):
     if tz is not None:
         monkeypatch.setenv("TZ", f"/usr/share/zoneinfo/{tz}")
+    tzset()
 
     # Roundtrip the time
     datetime2 = m.test_chrono2(time1)
@@ -263,6 +277,37 @@ def test_chrono_misc():
     assert roundtrip_datetime(d2) == d2
 
 
+
+@SKIP_TZ_ENV_ON_WIN
+@SKIP_TZ_NO_ZONEINFO
+@pytest.mark.parametrize(
+    "tz",
+    [
+        None,
+        pytest.param("UTC"),
+        pytest.param("America/New_York"),
+    ],
+)
+@pytest.mark.parametrize(
+    "sys_tz",
+    [
+        None,
+        pytest.param("UTC"),
+        pytest.param("America/New_York"),
+    ]
+)
+def test_chrono_tz_aware(sys_tz, tz, monkeypatch):
+    if sys_tz is not None:
+        monkeypatch.setenv("TZ", f"/usr/share/zoneinfo/{sys_tz}")
+    tzset()
+    from datetime import datetime, timedelta, timezone
+    from zoneinfo import ZoneInfo
+    roundtrip_datetime = m.test_nano_timepoint_roundtrip
+    now = datetime.now()
+    now_tz = now.astimezone(ZoneInfo(tz) if tz is not None else None)
+    assert roundtrip_datetime(now) == roundtrip_datetime(now_tz)
+
+
 @pytest.mark.parametrize(
     "test_type,roundtrip_name",
     [