add tests

Author: danielhochman

pysrc/orjson/__init__.py

@@ -11,6 +11,7 @@
     "JSONEncodeError",
     "loads",
     "OPT_APPEND_NEWLINE",
+    "OPT_DISALLOW_NAN",
     "OPT_INDENT_2",
     "OPT_NAIVE_UTC",
     "OPT_NON_STR_KEYS",

src/serialize/per_type/dict.rs

@@ -382,10 +382,17 @@ fn non_str_uuid(key: *mut pyo3_ffi::PyObject) -> Result<CompactString, Serialize
 
 #[cold]
 #[inline(never)]
-fn non_str_float(key: *mut pyo3_ffi::PyObject) -> Result<CompactString, SerializeError> {
+fn non_str_float(
+    key: *mut pyo3_ffi::PyObject,
+    opts: crate::opt::Opt,
+) -> Result<CompactString, SerializeError> {
     let val = ffi!(PyFloat_AS_DOUBLE(key));
     if !val.is_finite() {
-        Ok(CompactString::const_new("null"))
+        if unlikely!(opt_enabled!(opts, crate::opt::DISALLOW_NAN)) {
+            Err(SerializeError::FloatNotFinite)
+        } else {
+            Ok(CompactString::const_new("null"))
+        }
     } else {
         Ok(CompactString::from(ryu::Buffer::new().format_finite(val)))
     }
@@ -432,7 +439,7 @@ impl DictNonStrKey {
                 }
             }
             ObType::Int => non_str_int(key),
-            ObType::Float => non_str_float(key),
+            ObType::Float => non_str_float(key, opts),
             ObType::Datetime => non_str_datetime(key, opts),
             ObType::Date => non_str_date(key),
             ObType::Time => non_str_time(key, opts),

test/test_type.py

@@ -373,6 +373,46 @@ def test_infinity_loads(self):
         assert str(orjson.loads("[-Infinity]")[0]) == "-inf"
         assert str(orjson.loads("[-infinity]")[0]) == "-inf"
 
+    def test_disallow_nan_raises(self):
+        """
+        OPT_DISALLOW_NAN raises on NaN and Infinity
+        """
+        for val in (float("nan"), float("inf"), float("-inf")):
+            with pytest.raises(orjson.JSONEncodeError):
+                orjson.dumps(val, option=orjson.OPT_DISALLOW_NAN)
+
+    def test_disallow_nan_finite_ok(self):
+        """
+        OPT_DISALLOW_NAN allows finite floats
+        """
+        assert orjson.dumps(3.14, option=orjson.OPT_DISALLOW_NAN) == b"3.14"
+        assert orjson.dumps(0.0, option=orjson.OPT_DISALLOW_NAN) == b"0.0"
+        assert orjson.dumps(-1.5, option=orjson.OPT_DISALLOW_NAN) == b"-1.5"
+
+    def test_disallow_nan_in_list(self):
+        """
+        OPT_DISALLOW_NAN raises on NaN inside a list
+        """
+        with pytest.raises(orjson.JSONEncodeError):
+            orjson.dumps([1.0, float("nan"), 3.0], option=orjson.OPT_DISALLOW_NAN)
+
+    def test_disallow_nan_in_dict(self):
+        """
+        OPT_DISALLOW_NAN raises on Infinity inside a dict value
+        """
+        with pytest.raises(orjson.JSONEncodeError):
+            orjson.dumps({"x": float("inf")}, option=orjson.OPT_DISALLOW_NAN)
+
+    def test_disallow_nan_dict_key(self):
+        """
+        OPT_DISALLOW_NAN raises on NaN as a dict key with OPT_NON_STR_KEYS
+        """
+        with pytest.raises(orjson.JSONEncodeError):
+            orjson.dumps(
+                {float("nan"): 1},
+                option=orjson.OPT_DISALLOW_NAN | orjson.OPT_NON_STR_KEYS,
+            )
+
     def test_int_53(self):
         """
         int 53-bit