add an option to disallow nan

Author: danielhochman

pysrc/orjson/__init__.pyi

@@ -17,6 +17,7 @@ class Fragment(tuple):
     contents: Union[bytes, str]
 
 OPT_APPEND_NEWLINE: int
+OPT_DISALLOW_NAN: int
 OPT_INDENT_2: int
 OPT_NAIVE_UTC: int
 OPT_NON_STR_KEYS: int

src/lib.rs

@@ -126,6 +126,7 @@ pub unsafe extern "C" fn orjson_init_exec(mptr: *mut PyObject) -> c_int {
         add!(mptr, "Fragment\0", typeref::FRAGMENT_TYPE as *mut PyObject);
 
         opt!(mptr, "OPT_APPEND_NEWLINE\0", opt::APPEND_NEWLINE);
+        opt!(mptr, "OPT_DISALLOW_NAN\0", opt::DISALLOW_NAN);
         opt!(mptr, "OPT_INDENT_2\0", opt::INDENT_2);
         opt!(mptr, "OPT_NAIVE_UTC\0", opt::NAIVE_UTC);
         opt!(mptr, "OPT_NON_STR_KEYS\0", opt::NON_STR_KEYS);

src/opt.rs

@@ -14,6 +14,7 @@ pub const PASSTHROUGH_SUBCLASS: Opt = 1 << 8;
 pub const PASSTHROUGH_DATETIME: Opt = 1 << 9;
 pub const APPEND_NEWLINE: Opt = 1 << 10;
 pub const PASSTHROUGH_DATACLASS: Opt = 1 << 11;
+pub const DISALLOW_NAN: Opt = 1 << 12;
 
 // deprecated
 pub const SERIALIZE_DATACLASS: Opt = 0;
@@ -25,6 +26,7 @@ pub const NOT_PASSTHROUGH: Opt =
     !(PASSTHROUGH_DATETIME | PASSTHROUGH_DATACLASS | PASSTHROUGH_SUBCLASS);
 
 pub const MAX_OPT: i32 = (APPEND_NEWLINE
+    | DISALLOW_NAN
     | INDENT_2
     | NAIVE_UTC
     | NON_STR_KEYS

src/serialize/error.rs

@@ -6,6 +6,7 @@ use core::ptr::NonNull;
 pub enum SerializeError {
     DatetimeLibraryUnsupported,
     DefaultRecursionLimit,
+    FloatNotFinite,
     Integer53Bits,
     Integer64Bits,
     InvalidStr,
@@ -35,6 +36,9 @@ impl std::fmt::Display for SerializeError {
             SerializeError::DefaultRecursionLimit => {
                 write!(f, "default serializer exceeds recursion limit")
             }
+            SerializeError::FloatNotFinite => {
+                write!(f, "Cannot serialize Infinity or NaN")
+            }
             SerializeError::Integer53Bits => write!(f, "Integer exceeds 53-bit range"),
             SerializeError::Integer64Bits => write!(f, "Integer exceeds 64-bit range"),
             SerializeError::InvalidStr => write!(f, "{}", crate::util::INVALID_STR),

src/serialize/per_type/dict.rs

@@ -110,7 +110,7 @@ macro_rules! impl_serialize_entry {
             }
             ObType::Float => {
                 $map.serialize_key($key).unwrap();
-                $map.serialize_value(&FloatSerializer::new($value))?;
+                $map.serialize_value(&FloatSerializer::new($value, $self.state.opts()))?;
             }
             ObType::Bool => {
                 $map.serialize_key($key).unwrap();

src/serialize/per_type/float.rs

@@ -1,15 +1,17 @@
 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+use crate::opt::{Opt, DISALLOW_NAN};
+use crate::serialize::error::SerializeError;
 use serde::ser::{Serialize, Serializer};
 
-#[repr(transparent)]
 pub struct FloatSerializer {
     ptr: *mut pyo3_ffi::PyObject,
+    opts: Opt,
 }
 
 impl FloatSerializer {
-    pub fn new(ptr: *mut pyo3_ffi::PyObject) -> Self {
-        FloatSerializer { ptr: ptr }
+    pub fn new(ptr: *mut pyo3_ffi::PyObject, opts: Opt) -> Self {
+        FloatSerializer { ptr: ptr, opts: opts }
     }
 }
 
@@ -22,6 +24,9 @@ impl Serialize for FloatSerializer {
         let value = ffi!(PyFloat_AS_DOUBLE(self.ptr));
         #[cfg(yyjson_allow_inf_and_nan)]
         {
+            if unlikely!(opt_enabled!(self.opts, DISALLOW_NAN)) && !value.is_finite() {
+                err!(SerializeError::FloatNotFinite)
+            }
             serializer.serialize_f64(value)
         }
         #[cfg(not(yyjson_allow_inf_and_nan))]

src/serialize/per_type/list.rs

@@ -106,7 +106,7 @@ impl Serialize for ListTupleSerializer {
                     seq.serialize_element(&NoneSerializer::new()).unwrap();
                 }
                 ObType::Float => {
-                    seq.serialize_element(&FloatSerializer::new(value))?;
+                    seq.serialize_element(&FloatSerializer::new(value, self.state.opts()))?;
                 }
                 ObType::Bool => {
                     seq.serialize_element(&BoolSerializer::new(value)).unwrap();

src/serialize/serializer.rs

@@ -70,7 +70,7 @@ impl Serialize for PyObjectSerializer {
             ObType::StrSubclass => StrSubclassSerializer::new(self.ptr).serialize(serializer),
             ObType::Int => IntSerializer::new(self.ptr, self.state.opts()).serialize(serializer),
             ObType::None => NoneSerializer::new().serialize(serializer),
-            ObType::Float => FloatSerializer::new(self.ptr).serialize(serializer),
+            ObType::Float => FloatSerializer::new(self.ptr, self.state.opts()).serialize(serializer),
             ObType::Bool => BoolSerializer::new(self.ptr).serialize(serializer),
             ObType::Datetime => DateTime::new(self.ptr, self.state.opts()).serialize(serializer),
             ObType::Date => Date::new(self.ptr).serialize(serializer),