Fixes for ser_json_inf_nan inference (#1251)

Author: sydney-runkle

src/serializers/infer.rs

@@ -401,7 +401,14 @@ pub(crate) fn infer_serialize_known<S: Serializer>(
         ObType::None => serializer.serialize_none(),
         ObType::Int | ObType::IntSubclass => serialize!(Int),
         ObType::Bool => serialize!(bool),
-        ObType::Float | ObType::FloatSubclass => serialize!(f64),
+        ObType::Float | ObType::FloatSubclass => {
+            let v = value.extract::<f64>().map_err(py_err_se_err)?;
+            if (v.is_nan() || v.is_infinite()) && extra.config.inf_nan_mode == InfNanMode::Null {
+                serializer.serialize_none()
+            } else {
+                serializer.serialize_f64(v)
+            }
+        }
         ObType::Decimal => value.to_string().serialize(serializer),
         ObType::Str | ObType::StrSubclass => {
             let py_str = value.downcast::<PyString>().map_err(py_err_se_err)?;
@@ -575,8 +582,14 @@ pub(crate) fn infer_json_key_known<'a>(
 ) -> PyResult<Cow<'a, str>> {
     match ob_type {
         ObType::None => super::type_serializers::simple::none_json_key(),
-        ObType::Int | ObType::IntSubclass | ObType::Float | ObType::FloatSubclass => {
-            super::type_serializers::simple::to_str_json_key(key)
+        ObType::Int | ObType::IntSubclass => super::type_serializers::simple::to_str_json_key(key),
+        ObType::Float | ObType::FloatSubclass => {
+            let v = key.extract::<f64>()?;
+            if (v.is_nan() || v.is_infinite()) && extra.config.inf_nan_mode == InfNanMode::Null {
+                super::type_serializers::simple::none_json_key()
+            } else {
+                super::type_serializers::simple::to_str_json_key(key)
+            }
         }
         ObType::Decimal => Ok(Cow::Owned(key.to_string())),
         ObType::Bool => super::type_serializers::simple::bool_json_key(key),

tests/serializers/test_any.py

@@ -7,6 +7,7 @@
 from datetime import date, datetime, time, timedelta, timezone
 from decimal import Decimal
 from enum import Enum
+from math import inf, isinf, isnan, nan
 from pathlib import Path
 from typing import ClassVar
 
@@ -603,3 +604,43 @@ def test_numpy_float(any_serializer):
         any_serializer.to_python(numpy.float16(1.0), mode='json')
     with pytest.raises(PydanticSerializationError, match=r"Unable to serialize unknown type: <class 'numpy\.float16'>"):
         any_serializer.to_json(numpy.float16(1.0))
+
+
+def test_ser_json_inf_nan_with_any() -> None:
+    s = SchemaSerializer(core_schema.any_schema(), core_schema.CoreConfig(ser_json_inf_nan='constants'))
+    assert isinf(s.to_python(inf))
+    assert isinf(s.to_python(inf, mode='json'))
+    assert s.to_json(inf) == b'Infinity'
+    assert isnan(s.to_python(nan))
+    assert isnan(s.to_python(nan, mode='json'))
+    assert s.to_json(nan) == b'NaN'
+
+    s = SchemaSerializer(core_schema.any_schema(), core_schema.CoreConfig(ser_json_inf_nan='null'))
+    assert isinf(s.to_python(inf))
+    assert s.to_python(inf, mode='json') is None
+    assert s.to_json(inf) == b'null'
+    assert isnan(s.to_python(nan))
+    assert s.to_python(nan, mode='json') is None
+    assert s.to_json(nan) == b'null'
+
+
+def test_ser_json_inf_nan_with_list_of_any() -> None:
+    s = SchemaSerializer(
+        core_schema.list_schema(core_schema.any_schema()), core_schema.CoreConfig(ser_json_inf_nan='constants')
+    )
+    assert isinf(s.to_python([inf])[0])
+    assert isinf(s.to_python([inf], mode='json')[0])
+    assert s.to_json([inf]) == b'[Infinity]'
+    assert isnan(s.to_python([nan])[0])
+    assert isnan(s.to_python([nan], mode='json')[0])
+    assert s.to_json([nan]) == b'[NaN]'
+
+    s = SchemaSerializer(
+        core_schema.list_schema(core_schema.any_schema()), core_schema.CoreConfig(ser_json_inf_nan='null')
+    )
+    assert isinf(s.to_python([inf])[0])
+    assert s.to_python([inf], mode='json')[0] is None
+    assert s.to_json([inf]) == b'[null]'
+    assert isnan(s.to_python([nan])[0])
+    assert s.to_python([nan], mode='json')[0] is None
+    assert s.to_json([nan]) == b'[null]'