Add PyTzInfoAccess

src/types/datetime.rs

@@ -5,6 +5,9 @@
 
 use crate::err::PyResult;
 use crate::ffi;
+#[cfg(PyPy)]
+use crate::ffi::datetime::{PyDateTime_FromTimestamp, PyDate_FromTimestamp};
+use crate::ffi::PyDateTime_DATE_GET_TZINFO;
 use crate::ffi::{
     PyDateTime_CAPI, PyDateTime_FromTimestamp, PyDateTime_IMPORT, PyDate_FromTimestamp,
 };
@@ -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,11 @@ pub trait PyTimeAccess {
     fn get_fold(&self) -> bool;
 }
 
+/// Trait for accessing the components of a struct containing a tzinfo.
+pub trait PyTzInfoAccess {
+    fn get_tzinfo(&self) -> Option<&PyTzInfo>;
+}
+
 /// Bindings around `datetime.date`
 #[repr(transparent)]
 pub struct PyDate(PyAny);
@@ -354,6 +363,15 @@ impl PyTimeAccess for PyDateTime {
     }
 }
 
+impl PyTzInfoAccess for PyDateTime {
+    fn get_tzinfo(&self) -> Option<&PyTzInfo> {
+        unsafe {
+            self.py()
+                .from_owned_ptr_or_opt(PyDateTime_DATE_GET_TZINFO(self.as_ptr()))
+        }
+    }
+}
+
 /// Bindings for `datetime.time`
 #[repr(transparent)]
 pub struct PyTime(PyAny);
@@ -524,4 +542,25 @@ 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, 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());
+        });
+    }
 }

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;