Fast conditional for zoneinfo.ZoneInfo

ijl · Jan 21, 2022 · 8fc1e89 · 8fc1e89
1 parent 853ffbd
commit 8fc1e89
Show file tree

Hide file tree

Showing 9 changed files with 213 additions and 64 deletions.
diff --git a/README.md b/README.md
@@ -557,13 +557,13 @@ Serialize a UTC timezone on `datetime.datetime` instances as `Z` instead
 of `+00:00`.
 
 ```python
->>> import orjson, datetime
+>>> import orjson, datetime, zoneinfo
 >>> orjson.dumps(
-        datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
+        datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=zoneinfo.ZoneInfo("UTC")),
     )
 b'"1970-01-01T00:00:00+00:00"'
 >>> orjson.dumps(
-        datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
+        datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=zoneinfo.ZoneInfo("UTC")),
         option=orjson.OPT_UTC_Z
     )
 b'"1970-01-01T00:00:00Z"'
@@ -653,21 +653,21 @@ e.g., `field(metadata={"json_serialize": False})`, if use cases are clear.
 
 orjson serializes `datetime.datetime` objects to
 [RFC 3339](https://tools.ietf.org/html/rfc3339) format,
-e.g., "1970-01-01T00:00:00+00:00". This is a subset of ISO 8601 and
+e.g., "1970-01-01T00:00:00+00:00". This is a subset of ISO 8601 and is
 compatible with `isoformat()` in the standard library.
 
 ```python
 >>> import orjson, datetime, zoneinfo
 >>> orjson.dumps(
-    datetime.datetime(2018, 12, 1, 2, 3, 4, 9, tzinfo=zoneinfo.ZoneInfo('Australia/Adelaide'))
+    datetime.datetime(2018, 12, 1, 2, 3, 4, 9, tzinfo=zoneinfo.ZoneInfo("Australia/Adelaide"))
 )
 b'"2018-12-01T02:03:04.000009+10:30"'
 >>> orjson.dumps(
-    datetime.datetime.fromtimestamp(4123518902).replace(tzinfo=datetime.timezone.utc)
+    datetime.datetime(2100, 9, 1, 21, 55, 2).replace(tzinfo=zoneinfo.ZoneInfo("UTC"))
 )
 b'"2100-09-01T21:55:02+00:00"'
 >>> orjson.dumps(
-    datetime.datetime.fromtimestamp(4123518902)
+    datetime.datetime(2100, 9, 1, 21, 55, 2)
 )
 b'"2100-09-01T21:55:02"'
 ```
@@ -677,6 +677,8 @@ b'"2100-09-01T21:55:02"'
 module, or a timezone instance from the third-party `pendulum`, `pytz`, or
 `dateutil`/`arrow` libraries.
 
+It is fastest to use the standard library's `zoneinfo.ZoneInfo` for timezones.
+
 `datetime.time` objects must not have a `tzinfo`.
 
 ```python

diff --git a/src/serialize/datetime.rs b/src/serialize/datetime.rs
@@ -162,33 +162,50 @@ impl DateTimeLike for DateTime {
         unsafe { (*(self.ptr as *mut pyo3::ffi::PyDateTime_DateTime)).hastzinfo == 1 }
     }
 
+    fn slow_offset(&self) -> Result<Offset, DateTimeError> {
+        let tzinfo = ffi!(PyDateTime_DATE_GET_TZINFO(self.ptr));
+        if ffi!(PyObject_HasAttr(tzinfo, CONVERT_METHOD_STR)) == 1 {
+            // pendulum
+            let py_offset = call_method!(self.ptr, UTCOFFSET_METHOD_STR);
+            let offset = Offset {
+                second: ffi!(PyDateTime_DELTA_GET_SECONDS(py_offset)) as i32,
+                day: ffi!(PyDateTime_DELTA_GET_DAYS(py_offset)),
+            };
+            ffi!(Py_DECREF(py_offset));
+            Ok(offset)
+        } else if ffi!(PyObject_HasAttr(tzinfo, NORMALIZE_METHOD_STR)) == 1 {
+            // pytz
+            let method_ptr = call_method!(tzinfo, NORMALIZE_METHOD_STR, self.ptr);
+            let py_offset = call_method!(method_ptr, UTCOFFSET_METHOD_STR);
+            ffi!(Py_DECREF(method_ptr));
+            let offset = Offset {
+                second: ffi!(PyDateTime_DELTA_GET_SECONDS(py_offset)) as i32,
+                day: ffi!(PyDateTime_DELTA_GET_DAYS(py_offset)),
+            };
+            ffi!(Py_DECREF(py_offset));
+            Ok(offset)
+        } else if ffi!(PyObject_HasAttr(tzinfo, DST_STR)) == 1 {
+            // dateutil/arrow, datetime.timezone.utc
+            let py_offset = call_method!(tzinfo, UTCOFFSET_METHOD_STR, self.ptr);
+            let offset = Offset {
+                second: ffi!(PyDateTime_DELTA_GET_SECONDS(py_offset)) as i32,
+                day: ffi!(PyDateTime_DELTA_GET_DAYS(py_offset)),
+            };
+            ffi!(Py_DECREF(py_offset));
+            Ok(offset)
+        } else {
+            Err(DateTimeError::LibraryUnsupported)
+        }
+    }
+
+    #[cfg(Py_3_9)]
     fn offset(&self) -> Result<Offset, DateTimeError> {
         if !self.has_tz() {
             Ok(Offset::default())
         } else {
             let tzinfo = ffi!(PyDateTime_DATE_GET_TZINFO(self.ptr));
-            if ffi!(PyObject_HasAttr(tzinfo, CONVERT_METHOD_STR)) == 1 {
-                // pendulum
-                let py_offset = call_method!(self.ptr, UTCOFFSET_METHOD_STR);
-                let offset = Offset {
-                    second: ffi!(PyDateTime_DELTA_GET_SECONDS(py_offset)) as i32,
-                    day: ffi!(PyDateTime_DELTA_GET_DAYS(py_offset)),
-                };
-                ffi!(Py_DECREF(py_offset));
-                Ok(offset)
-            } else if ffi!(PyObject_HasAttr(tzinfo, NORMALIZE_METHOD_STR)) == 1 {
-                // pytz
-                let method_ptr = call_method!(tzinfo, NORMALIZE_METHOD_STR, self.ptr);
-                let py_offset = call_method!(method_ptr, UTCOFFSET_METHOD_STR);
-                ffi!(Py_DECREF(method_ptr));
-                let offset = Offset {
-                    second: ffi!(PyDateTime_DELTA_GET_SECONDS(py_offset)) as i32,
-                    day: ffi!(PyDateTime_DELTA_GET_DAYS(py_offset)),
-                };
-                ffi!(Py_DECREF(py_offset));
-                Ok(offset)
-            } else if ffi!(PyObject_HasAttr(tzinfo, DST_STR)) == 1 {
-                // dateutil/arrow, datetime.timezone.utc
+            if unsafe { ob_type!(tzinfo) == ZONEINFO_TYPE } {
+                // zoneinfo
                 let py_offset = call_method!(tzinfo, UTCOFFSET_METHOD_STR, self.ptr);
                 let offset = Offset {
                     second: ffi!(PyDateTime_DELTA_GET_SECONDS(py_offset)) as i32,
@@ -197,10 +214,19 @@ impl DateTimeLike for DateTime {
                 ffi!(Py_DECREF(py_offset));
                 Ok(offset)
             } else {
-                Err(DateTimeError::LibraryUnsupported)
+                self.slow_offset()
             }
         }
     }
+
+    #[cfg(not(Py_3_9))]
+    fn offset(&self) -> Result<Offset, DateTimeError> {
+        if !self.has_tz() {
+            Ok(Offset::default())
+        } else {
+            self.slow_offset()
+        }
+    }
 }
 
 impl<'p> Serialize for DateTime {

diff --git a/src/serialize/datetimelike.rs b/src/serialize/datetimelike.rs
@@ -87,6 +87,9 @@ pub trait DateTimeLike {
     /// Is the object time-zone aware?
     fn has_tz(&self) -> bool;
 
+    //// python3.8 or below implementation of offset()
+    fn slow_offset(&self) -> Result<Offset, DateTimeError>;
+
     /// The offset of the timezone.
     fn offset(&self) -> Result<Offset, DateTimeError>;
 
@@ -116,15 +119,9 @@ pub trait DateTimeLike {
         if opts & OMIT_MICROSECONDS == 0 {
             let microsecond = self.microsecond();
             if microsecond != 0 {
-                {
-                    buf.push(b'.');
-                }
-                {
-                    write_triple_digit!(buf, microsecond / 1_000);
-                }
-                {
-                    write_triple_digit!(buf, microsecond % 1_000);
-                }
+                buf.push(b'.');
+                write_triple_digit!(buf, microsecond / 1_000);
+                write_triple_digit!(buf, microsecond % 1_000);
                 // Don't support writing nanoseconds for now.
                 // If requested, something like the following should work,
                 // and the `DateTimeBuffer` type alias should be changed to
@@ -155,23 +152,21 @@ pub trait DateTimeLike {
                     // datetime.timedelta(seconds=37800) -> +10:30
                     buf.push(b'+');
                 }
-                {
-                    let offset_minute = offset_second / 60;
-                    let offset_hour = offset_minute / 60;
-                    write_double_digit!(buf, offset_hour);
-                    buf.push(b':');
-                    let mut offset_minute_print = offset_minute % 60;
-                    // https://tools.ietf.org/html/rfc3339#section-5.8
-                    // "exactly 19 minutes and 32.13 seconds ahead of UTC"
-                    // "closest representable UTC offset"
-                    //  "+20:00"
-                    let offset_excess_second =
-                        offset_second - (offset_minute_print * 60 + offset_hour * 3600);
-                    if offset_excess_second >= 30 {
-                        offset_minute_print += 1;
-                    }
-                    write_double_digit!(buf, offset_minute_print);
+                let offset_minute = offset_second / 60;
+                let offset_hour = offset_minute / 60;
+                write_double_digit!(buf, offset_hour);
+                buf.push(b':');
+                let mut offset_minute_print = offset_minute % 60;
+                // https://tools.ietf.org/html/rfc3339#section-5.8
+                // "exactly 19 minutes and 32.13 seconds ahead of UTC"
+                // "closest representable UTC offset"
+                //  "+20:00"
+                let offset_excess_second =
+                    offset_second - (offset_minute_print * 60 + offset_hour * 3600);
+                if offset_excess_second >= 30 {
+                    offset_minute_print += 1;
                 }
+                write_double_digit!(buf, offset_minute_print);
             }
         }
         Ok(())

diff --git a/src/serialize/numpy.rs b/src/serialize/numpy.rs
@@ -837,6 +837,10 @@ impl DateTimeLike for NumpyDatetime64Repr {
         false
     }
 
+    fn slow_offset(&self) -> Result<Offset, DateTimeError> {
+        unreachable!()
+    }
+
     fn offset(&self) -> Result<Offset, DateTimeError> {
         Ok(Offset::default())
     }

diff --git a/src/typeref.rs b/src/typeref.rs
@@ -45,6 +45,9 @@ pub static mut TUPLE_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
 pub static mut UUID_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
 pub static mut ENUM_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
 
+#[cfg(Py_3_9)]
+pub static mut ZONEINFO_TYPE: *mut PyTypeObject = 0 as *mut PyTypeObject;
+
 pub static mut NUMPY_TYPES: Lazy<Option<NumpyTypes>> = Lazy::new(|| unsafe { load_numpy_types() });
 pub static mut FIELD_TYPE: Lazy<NonNull<PyObject>> = Lazy::new(|| unsafe { look_up_field_type() });
 
@@ -120,6 +123,11 @@ pub fn init_typerefs() {
         UUID_TYPE = look_up_uuid_type();
         ENUM_TYPE = look_up_enum_type();
 
+        #[cfg(Py_3_9)]
+        {
+            ZONEINFO_TYPE = look_up_zoneinfo_type();
+        }
+
         INT_ATTR_STR = PyUnicode_InternFromString("int\0".as_ptr() as *const c_char);
         UTCOFFSET_METHOD_STR = PyUnicode_InternFromString("utcoffset\0".as_ptr() as *const c_char);
         NORMALIZE_METHOD_STR = PyUnicode_InternFromString("normalize\0".as_ptr() as *const c_char);
@@ -269,3 +277,16 @@ unsafe fn look_up_time_type() -> *mut PyTypeObject {
     Py_DECREF(time);
     ptr
 }
+
+#[cfg(Py_3_9)]
+#[cold]
+#[cfg_attr(feature = "unstable-simd", optimize(size))]
+unsafe fn look_up_zoneinfo_type() -> *mut PyTypeObject {
+    let module = PyImport_ImportModule("zoneinfo\0".as_ptr() as *const c_char);
+    let module_dict = PyObject_GenericGetDict(module, std::ptr::null_mut());
+    let ptr = PyMapping_GetItemString(module_dict, "ZoneInfo\0".as_ptr() as *const c_char)
+        as *mut PyTypeObject;
+    Py_DECREF(module_dict);
+    Py_DECREF(module);
+    ptr
+}