From e95e73b55e2f3d8329f1d5720cc83a36a0da2194 Mon Sep 17 00:00:00 2001
From: David Hewitt <1939362+davidhewitt@users.noreply.github.com>
Date: Wed, 21 Apr 2021 21:26:47 +0100
Subject: [PATCH] ffi: support PyDateTime_TimeZone_UTC

---
 CHANGELOG.md        |  1 +
 src/ffi/datetime.rs | 91 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 52b18938..ae1222ed 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - Add FFI definition `Py_IS_TYPE`. [#1429](https://github.com/PyO3/pyo3/pull/1429)
 - Add FFI definition `_Py_InitializeMain`. [#1473](https://github.com/PyO3/pyo3/pull/1473)
 - Add tuple and unit struct support for `#[pyclass]` macro. [#1504](https://github.com/PyO3/pyo3/pull/1504)
+- Add FFI definition `PyDateTime_TimeZone_UTC`. [#1572](https://github.com/PyO3/pyo3/pull/1572)
 
 ### Changed
 - Change `PyTimeAcces::get_fold()` to return a `bool` instead of a `u8`. [#1397](https://github.com/PyO3/pyo3/pull/1397)
diff --git a/src/ffi/datetime.rs b/src/ffi/datetime.rs
index 70d8f104..5f4915df 100644
--- a/src/ffi/datetime.rs
+++ b/src/ffi/datetime.rs
@@ -399,13 +399,11 @@ pub struct PyDateTime_CAPI {
     pub TimeZone_FromTimeZone:
         unsafe extern "C" fn(offset: *mut PyObject, name: *mut PyObject) -> *mut PyObject,
 
-    // Defined for PyPy as `PyDateTime_FromTimestamp`
     pub DateTime_FromTimestamp: unsafe extern "C" fn(
         cls: *mut PyTypeObject,
         args: *mut PyObject,
         kwargs: *mut PyObject,
     ) -> *mut PyObject,
-    // Defined for PyPy as `PyDate_FromTimestamp`
     pub Date_FromTimestamp:
         unsafe extern "C" fn(cls: *mut PyTypeObject, args: *mut PyObject) -> *mut PyObject,
     #[cfg(not(PyPy))]
@@ -447,6 +445,17 @@ pub static PyDateTimeAPI: _PyDateTimeAPI_impl = _PyDateTimeAPI_impl {
     inner: GILOnceCell::new(),
 };
 
+/// Safe wrapper around the Python C-API global `PyDateTime_TimeZone_UTC`. This follows a similar
+/// strategy as [`PyDateTimeAPI`]: the Python datetime C-API will automatically be imported if this
+/// type is deferenced.
+///
+/// The type obtained by dereferencing this object is `&'static PyObject`. This may change in the
+/// future to be a more specific type representing that this is a `datetime.timezone` object.
+#[cfg(all(Py_3_7, not(PyPy)))]
+pub static PyDateTime_TimeZone_UTC: _PyDateTime_TimeZone_UTC_impl = _PyDateTime_TimeZone_UTC_impl {
+    inner: &PyDateTimeAPI,
+};
+
 /// Populates the `PyDateTimeAPI` object
 ///
 /// Unlike in C, this does *not* need to be actively invoked in Rust, which
@@ -559,6 +568,16 @@ pub unsafe fn PyTZInfo_CheckExact(op: *mut PyObject) -> c_int {
 // skipped non-limited PyTimeZone_FromOffset
 // skipped non-limited PyTimeZone_FromOffsetAndName
 
+#[cfg(not(PyPy))]
+pub unsafe fn PyDateTime_FromTimestamp(args: *mut PyObject) -> *mut PyObject {
+    (PyDateTimeAPI.DateTime_FromTimestamp)(PyDateTimeAPI.DateTimeType, args, std::ptr::null_mut())
+}
+
+#[cfg(not(PyPy))]
+pub unsafe fn PyDate_FromTimestamp(args: *mut PyObject) -> *mut PyObject {
+    (PyDateTimeAPI.Date_FromTimestamp)(PyDateTimeAPI.DateType, args)
+}
+
 #[cfg(PyPy)]
 extern "C" {
     #[link_name = "PyPyDate_FromTimestamp"]
@@ -566,6 +585,7 @@ extern "C" {
     #[link_name = "PyPyDateTime_FromTimestamp"]
     pub fn PyDateTime_FromTimestamp(args: *mut PyObject) -> *mut PyObject;
 }
+
 #[cfg(PyPy)]
 extern "C" {
     #[link_name = "_PyPyDateTime_Import"]
@@ -587,3 +607,70 @@ impl Deref for _PyDateTimeAPI_impl {
         unsafe { PyDateTime_IMPORT() }
     }
 }
+
+#[doc(hidden)]
+#[cfg(all(Py_3_7, not(PyPy)))]
+pub struct _PyDateTime_TimeZone_UTC_impl {
+    inner: &'static _PyDateTimeAPI_impl,
+}
+
+#[cfg(all(Py_3_7, not(PyPy)))]
+impl Deref for _PyDateTime_TimeZone_UTC_impl {
+    type Target = crate::PyObject;
+
+    #[inline]
+    fn deref(&self) -> &crate::PyObject {
+        unsafe {
+            &*((&self.inner.TimeZone_UTC) as *const *mut crate::ffi::PyObject
+                as *const crate::PyObject)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{py_run, AsPyPointer, IntoPy, Py, PyAny, Python};
+
+    #[test]
+    fn test_datetime_fromtimestamp() {
+        Python::with_gil(|py| {
+            let args: Py<PyAny> = (100,).into_py(py);
+            unsafe { PyDateTime_IMPORT() };
+            let dt: &PyAny = unsafe { py.from_owned_ptr(PyDateTime_FromTimestamp(args.as_ptr())) };
+            py_run!(
+                py,
+                dt,
+                "import datetime; assert dt == datetime.datetime.fromtimestamp(100)"
+            );
+        })
+    }
+
+    #[test]
+    fn test_date_fromtimestamp() {
+        Python::with_gil(|py| {
+            let args: Py<PyAny> = (100,).into_py(py);
+            dbg!(args.as_ref(py));
+            unsafe { PyDateTime_IMPORT() };
+            let dt: &PyAny = unsafe { py.from_owned_ptr(PyDate_FromTimestamp(args.as_ptr())) };
+            py_run!(
+                py,
+                dt,
+                "import datetime; assert dt == datetime.date.fromtimestamp(100)"
+            );
+        })
+    }
+
+    #[test]
+    #[cfg(all(Py_3_7, not(PyPy)))]
+    fn test_utc_timezone() {
+        Python::with_gil(|py| {
+            let utc_timezone = PyDateTime_TimeZone_UTC.as_ref(py);
+            py_run!(
+                py,
+                utc_timezone,
+                "import datetime; assert utc_timezone is datetime.timezone.utc"
+            );
+        })
+    }
+}