Refactor type serializers

Author: ijl

src/datetime.rs

@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+use crate::exc::*;
 use crate::typeref::*;
-use serde::ser::{Error, Serialize, Serializer};
+use serde::ser::{Serialize, Serializer};
 
 pub const NAIVE_UTC: u16 = 1 << 1;
 pub const OMIT_MICROSECONDS: u16 = 1 << 2;
@@ -15,12 +16,6 @@ const COLON: u8 = 58; // ":"
 const PERIOD: u8 = 46; // ":"
 const Z: u8 = 90; // "Z"
 
-macro_rules! err {
-    ($msg:expr) => {
-        return Err(Error::custom($msg));
-    };
-}
-
 pub type DateTimeBuffer = heapless::Vec<u8, heapless::consts::U32>;
 
 macro_rules! write_double_digit {
@@ -55,13 +50,7 @@ impl Date {
     pub fn new(ptr: *mut pyo3::ffi::PyObject) -> Self {
         Date { ptr: ptr }
     }
-}
-impl<'p> Serialize for Date {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        let mut buf: DateTimeBuffer = heapless::Vec::new();
+    pub fn write_buf(&self, buf: &mut DateTimeBuffer) {
         {
             let year = ffi!(PyDateTime_GET_YEAR(self.ptr)) as i32;
             buf.extend_from_slice(itoa::Buffer::new().format(year).as_bytes())
@@ -77,33 +66,39 @@ impl<'p> Serialize for Date {
             let day = ffi!(PyDateTime_GET_DAY(self.ptr)) as u32;
             write_double_digit!(buf, day);
         }
+    }
+}
+impl<'p> Serialize for Date {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let mut buf: DateTimeBuffer = heapless::Vec::new();
+        self.write_buf(&mut buf);
         serializer.serialize_str(str_from_slice!(buf.as_ptr(), buf.len()))
     }
 }
 
+pub enum TimeError {
+    HasTimezone,
+}
+
 pub struct Time {
     ptr: *mut pyo3::ffi::PyObject,
     opts: u16,
 }
 
 impl Time {
-    pub fn new(ptr: *mut pyo3::ffi::PyObject, opts: u16) -> Self {
-        Time {
+    pub fn new(ptr: *mut pyo3::ffi::PyObject, opts: u16) -> Result<Self, TimeError> {
+        if unsafe { (*(ptr as *mut pyo3::ffi::PyDateTime_Time)).hastzinfo == 1 } {
+            return Err(TimeError::HasTimezone);
+        }
+        Ok(Time {
             ptr: ptr,
             opts: opts,
-        }
+        })
     }
-}
-
-impl<'p> Serialize for Time {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        if unsafe { (*(self.ptr as *mut pyo3::ffi::PyDateTime_Time)).hastzinfo == 1 } {
-            err!("datetime.time must not have tzinfo set")
-        }
-        let mut buf: DateTimeBuffer = heapless::Vec::new();
+    pub fn write_buf(&self, buf: &mut DateTimeBuffer) {
         {
             let hour = ffi!(PyDateTime_TIME_GET_HOUR(self.ptr)) as u8;
             write_double_digit!(buf, hour);
@@ -122,10 +117,24 @@ impl<'p> Serialize for Time {
             let microsecond = ffi!(PyDateTime_TIME_GET_MICROSECOND(self.ptr)) as u32;
             write_microsecond!(buf, microsecond);
         }
+    }
+}
+
+impl<'p> Serialize for Time {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let mut buf: DateTimeBuffer = heapless::Vec::new();
+        self.write_buf(&mut buf);
         serializer.serialize_str(str_from_slice!(buf.as_ptr(), buf.len()))
     }
 }
 
+pub enum DateTimeError {
+    LibraryUnsupported,
+}
+
 pub struct DateTime {
     ptr: *mut pyo3::ffi::PyObject,
     opts: u16,
@@ -138,14 +147,7 @@ impl DateTime {
             opts: opts,
         }
     }
-}
-
-impl<'p> Serialize for DateTime {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        let mut buf: DateTimeBuffer = heapless::Vec::new();
+    pub fn write_buf(&self, buf: &mut DateTimeBuffer) -> Result<(), DateTimeError> {
         let has_tz = unsafe { (*(self.ptr as *mut pyo3::ffi::PyDateTime_DateTime)).hastzinfo == 1 };
         let offset_day: i32;
         let mut offset_second: i32;
@@ -195,7 +197,7 @@ impl<'p> Serialize for DateTime {
                     offset_second = ffi!(PyDateTime_DELTA_GET_SECONDS(offset)) as i32;
                     offset_day = ffi!(PyDateTime_DELTA_GET_DAYS(offset));
                 } else {
-                    err!("datetime's timezone library is not supported: use datetime.timezone.utc, pendulum, pytz, or dateutil")
+                    return Err(DateTimeError::LibraryUnsupported);
                 }
             } else {
                 offset_second = 0;
@@ -285,6 +287,19 @@ impl<'p> Serialize for DateTime {
                 }
             }
         }
+        Ok(())
+    }
+}
+
+impl<'p> Serialize for DateTime {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let mut buf: DateTimeBuffer = heapless::Vec::new();
+        if self.write_buf(&mut buf).is_err() {
+            err!(DATETIME_LIBRARY_UNSUPPORTED)
+        }
         serializer.serialize_str(str_from_slice!(buf.as_ptr(), buf.len()))
     }
 }

src/default.rs

@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+use crate::encode::*;
+
+use serde::ser::{Serialize, Serializer};
+use std::ffi::CStr;
+
+use std::ptr::NonNull;
+
+macro_rules! obj_name {
+    ($obj:ident) => {
+        unsafe { CStr::from_ptr((*$obj).tp_name).to_string_lossy() }
+    };
+}
+
+pub struct DefaultSerializer {
+    ptr: *mut pyo3::ffi::PyObject,
+    opts: u16,
+    default_calls: u8,
+    recursion: u8,
+    default: Option<NonNull<pyo3::ffi::PyObject>>,
+}
+
+impl DefaultSerializer {
+    pub fn new(
+        ptr: *mut pyo3::ffi::PyObject,
+        opts: u16,
+        default_calls: u8,
+        recursion: u8,
+        default: Option<NonNull<pyo3::ffi::PyObject>>,
+    ) -> Self {
+        DefaultSerializer {
+            ptr: ptr,
+            opts: opts,
+            default_calls: default_calls,
+            recursion: recursion,
+            default: default,
+        }
+    }
+}
+
+impl<'p> Serialize for DefaultSerializer {
+    #[inline(never)]
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        match self.default {
+            Some(callable) => {
+                if unlikely!(self.default_calls == RECURSION_LIMIT) {
+                    err!("default serializer exceeds recursion limit")
+                }
+                let obj_ptr = unsafe { (*self.ptr).ob_type };
+                let default_obj = unsafe {
+                    pyo3::ffi::PyObject_CallFunctionObjArgs(
+                        callable.as_ptr(),
+                        self.ptr,
+                        std::ptr::null_mut() as *mut pyo3::ffi::PyObject,
+                    )
+                };
+                if default_obj.is_null() {
+                    err!(format_args!(
+                        "Type is not JSON serializable: {}",
+                        obj_name!(obj_ptr)
+                    ))
+                } else if !ffi!(PyErr_Occurred()).is_null() {
+                    err!(format_args!(
+                        "Type raised exception in default function: {}",
+                        obj_name!(obj_ptr)
+                    ))
+                } else {
+                    let res = SerializePyObject::new(
+                        default_obj,
+                        None,
+                        self.opts,
+                        self.default_calls + 1,
+                        self.recursion,
+                        self.default,
+                    )
+                    .serialize(serializer);
+                    ffi!(Py_DECREF(default_obj));
+                    res
+                }
+            }
+            None => {
+                let obj_ptr = unsafe { (*self.ptr).ob_type };
+                err!(format_args!(
+                    "Type is not JSON serializable: {}",
+                    obj_name!(obj_ptr)
+                ))
+            }
+        }
+    }
+}

src/dict.rs

@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+use crate::encode::*;
+use crate::exc::*;
+use crate::typeref::*;
+use crate::unicode::*;
+use serde::ser::{Serialize, SerializeMap, Serializer};
+use smallvec::SmallVec;
+use std::ptr::NonNull;
+
+pub struct DictSortedKey {
+    ptr: *mut pyo3::ffi::PyObject,
+    opts: u16,
+    default_calls: u8,
+    recursion: u8,
+    default: Option<NonNull<pyo3::ffi::PyObject>>,
+}
+
+impl DictSortedKey {
+    pub fn new(
+        ptr: *mut pyo3::ffi::PyObject,
+        opts: u16,
+        default_calls: u8,
+        recursion: u8,
+        default: Option<NonNull<pyo3::ffi::PyObject>>,
+    ) -> Self {
+        DictSortedKey {
+            ptr: ptr,
+            opts: opts,
+            default_calls: default_calls,
+            recursion: recursion,
+            default: default,
+        }
+    }
+}
+
+impl<'p> Serialize for DictSortedKey {
+    #[inline(never)]
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let len = ffi!(PyDict_Size(self.ptr)) as usize;
+        if len == 0 {
+            serializer.serialize_map(Some(0)).unwrap().end()
+        } else {
+            let mut items: SmallVec<[(&str, *mut pyo3::ffi::PyObject); 8]> =
+                SmallVec::with_capacity(len);
+            let mut pos = 0isize;
+            let mut str_size: pyo3::ffi::Py_ssize_t = 0;
+            let mut key: *mut pyo3::ffi::PyObject = std::ptr::null_mut();
+            let mut value: *mut pyo3::ffi::PyObject = std::ptr::null_mut();
+            while unsafe { pyo3::ffi::PyDict_Next(self.ptr, &mut pos, &mut key, &mut value) != 0 } {
+                if unlikely!((*key).ob_type != STR_TYPE) {
+                    err!("Dict key must be str")
+                }
+                let data = read_utf8_from_str(key, &mut str_size);
+                if unlikely!(data.is_null()) {
+                    err!(INVALID_STR)
+                }
+                items.push((str_from_slice!(data, str_size), value));
+            }
+
+            items.sort_unstable_by(|a, b| a.0.cmp(b.0));
+
+            let mut map = serializer.serialize_map(None).unwrap();
+            for (key, val) in items.iter() {
+                map.serialize_entry(
+                    key,
+                    &SerializePyObject::new(
+                        *val,
+                        None,
+                        self.opts,
+                        self.default_calls,
+                        self.recursion + 1,
+                        self.default,
+                    ),
+                )?;
+            }
+            map.end()
+        }
+    }
+}