From 0efd02b423ed298827d916f93c9e82f2a4deac26 Mon Sep 17 00:00:00 2001
From: Ivan Tham <pickfire@riseup.net>
Date: Fri, 1 Apr 2022 00:01:43 +0800
Subject: [PATCH] Add PyTzInfoAccess

---
 CHANGELOG.md          |  4 +++
 src/types/datetime.rs | 61 +++++++++++++++++++++++++++++++++++++++++++
 src/types/mod.rs      |  1 +
 3 files changed, 66 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6bd41b4c2de..1b4651e9301 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added `PyTzInfoAccess`. [#2263](https://github.com/PyO3/pyo3/pull/2263)
+
 ### Changed
 
 - Default to "m" ABI tag when choosing `libpython` link name for CPython 3.7 on Unix. [#2288](https://github.com/PyO3/pyo3/pull/2288)
diff --git a/src/types/datetime.rs b/src/types/datetime.rs
index c405301118e..1d65362ec0d 100644
--- a/src/types/datetime.rs
+++ b/src/types/datetime.rs
@@ -5,6 +5,8 @@
 
 use crate::err::PyResult;
 use crate::ffi;
+#[cfg(PyPy)]
+use crate::ffi::datetime::{PyDateTime_FromTimestamp, PyDate_FromTimestamp};
 use crate::ffi::{
     PyDateTime_CAPI, PyDateTime_FromTimestamp, PyDateTime_IMPORT, PyDate_FromTimestamp,
 };
@@ -14,6 +16,7 @@ use crate::ffi::{
     PyDateTime_DATE_GET_HOUR, PyDateTime_DATE_GET_MICROSECOND, PyDateTime_DATE_GET_MINUTE,
     PyDateTime_DATE_GET_SECOND,
 };
+use crate::ffi::{PyDateTime_DATE_GET_TZINFO, PyDateTime_TIME_GET_TZINFO};
 use crate::ffi::{
     PyDateTime_DELTA_GET_DAYS, PyDateTime_DELTA_GET_MICROSECONDS, PyDateTime_DELTA_GET_SECONDS,
 };
@@ -22,6 +25,7 @@ use crate::ffi::{
     PyDateTime_TIME_GET_HOUR, PyDateTime_TIME_GET_MICROSECOND, PyDateTime_TIME_GET_MINUTE,
     PyDateTime_TIME_GET_SECOND,
 };
+use crate::instance::PyNativeType;
 use crate::types::PyTuple;
 use crate::{AsPyPointer, PyAny, PyObject, Python, ToPyObject};
 use std::os::raw::c_int;
@@ -160,6 +164,16 @@ pub trait PyTimeAccess {
     fn get_fold(&self) -> bool;
 }
 
+/// Trait for accessing the components of a struct containing a tzinfo.
+pub trait PyTzInfoAccess {
+    /// Returns the tzinfo (which may be None).
+    ///
+    /// Implementations should conform to the upstream documentation:
+    /// <https://docs.python.org/3/c-api/datetime.html#c.PyDateTime_DATE_GET_TZINFO>
+    /// <https://docs.python.org/3/c-api/datetime.html#c.PyDateTime_TIME_GET_TZINFO>
+    fn get_tzinfo(&self) -> Option<&PyTzInfo>;
+}
+
 /// Bindings around `datetime.date`
 #[repr(transparent)]
 pub struct PyDate(PyAny);
@@ -354,6 +368,15 @@ impl PyTimeAccess for PyDateTime {
     }
 }
 
+impl PyTzInfoAccess for PyDateTime {
+    fn get_tzinfo(&self) -> Option<&PyTzInfo> {
+        unsafe {
+            self.py()
+                .from_borrowed_ptr_or_opt(PyDateTime_DATE_GET_TZINFO(self.as_ptr()))
+        }
+    }
+}
+
 /// Bindings for `datetime.time`
 #[repr(transparent)]
 pub struct PyTime(PyAny);
@@ -439,6 +462,15 @@ impl PyTimeAccess for PyTime {
     }
 }
 
+impl PyTzInfoAccess for PyTime {
+    fn get_tzinfo(&self) -> Option<&PyTzInfo> {
+        unsafe {
+            self.py()
+                .from_borrowed_ptr_or_opt(PyDateTime_TIME_GET_TZINFO(self.as_ptr()))
+        }
+    }
+}
+
 /// Bindings for `datetime.tzinfo`
 ///
 /// This is an abstract base class and should not be constructed directly.
@@ -524,4 +556,33 @@ mod tests {
             assert!(b.unwrap().get_fold());
         });
     }
+
+    #[cfg(not(PyPy))]
+    #[test]
+    fn test_get_tzinfo() {
+        crate::Python::with_gil(|py| {
+            use crate::conversion::ToPyObject;
+            use crate::types::{PyDateTime, PyTime, PyTzInfoAccess};
+
+            let datetime = py.import("datetime").map_err(|e| e.print(py)).unwrap();
+            let timezone = datetime.getattr("timezone").unwrap();
+            let utc = timezone.getattr("utc").unwrap().to_object(py);
+
+            let dt = PyDateTime::new(py, 2018, 1, 1, 0, 0, 0, 0, Some(&utc)).unwrap();
+
+            assert!(dt.get_tzinfo().unwrap().eq(&utc).unwrap());
+
+            let dt = PyDateTime::new(py, 2018, 1, 1, 0, 0, 0, 0, None).unwrap();
+
+            assert!(dt.get_tzinfo().is_none());
+
+            let t = PyTime::new(py, 0, 0, 0, 0, Some(&utc)).unwrap();
+
+            assert!(t.get_tzinfo().unwrap().eq(&utc).unwrap());
+
+            let t = PyTime::new(py, 0, 0, 0, 0, None).unwrap();
+
+            assert!(t.get_tzinfo().is_none());
+        });
+    }
 }
diff --git a/src/types/mod.rs b/src/types/mod.rs
index 192fa80a8fa..ad79d062b9a 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -11,6 +11,7 @@ pub use self::complex::PyComplex;
 #[cfg(not(Py_LIMITED_API))]
 pub use self::datetime::{
     PyDate, PyDateAccess, PyDateTime, PyDelta, PyDeltaAccess, PyTime, PyTimeAccess, PyTzInfo,
+    PyTzInfoAccess,
 };
 pub use self::dict::{IntoPyDict, PyDict};
 pub use self::floatob::PyFloat;