From 0d45d8bf2421a021eb4ead45829106faf6a1fef9 Mon Sep 17 00:00:00 2001 From: Kangxing Xie Date: Sat, 24 Aug 2024 19:58:09 +1000 Subject: [PATCH 1/2] feat: Add support for serializing IntSubclass values larger than i64 when using model_dump(..., mode='json') --- src/serializers/infer.rs | 16 ++++++++++------ tests/serializers/test_any.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/src/serializers/infer.rs b/src/serializers/infer.rs index 032ede46b..18b21356e 100644 --- a/src/serializers/infer.rs +++ b/src/serializers/infer.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; +use num_bigint::BigInt; use pyo3::exceptions::PyTypeError; use pyo3::intern; use pyo3::prelude::*; @@ -44,7 +45,6 @@ pub(crate) fn infer_to_python_known( mut extra: &Extra, ) -> PyResult { let py = value.py(); - let mode = extra.mode; let mut guard = match extra.recursion_guard(value, INFER_DEF_REF_ID) { Ok(v) => v, @@ -110,16 +110,20 @@ pub(crate) fn infer_to_python_known( ); serializer.serializer.to_python(value, include, exclude, &extra) }; - let value = match extra.mode { SerMode::Json => match ob_type { // `bool` and `None` can't be subclasses, `ObType::Int`, `ObType::Float`, `ObType::Str` refer to exact types ObType::None | ObType::Bool | ObType::Int | ObType::Str => value.into_py(py), // have to do this to make sure subclasses of for example str are upcast to `str` - ObType::IntSubclass => match extract_i64(value) { - Some(v) => v.into_py(py), - None => return py_err!(PyTypeError; "expected int, got {}", safe_repr(value)), - }, + ObType::IntSubclass => { + if let Some(v) = extract_i64(value) { + v.into_py(py) + } else if let Ok(b) = value.extract::() { + b.into_py(py) + } else { + return py_err!(PyTypeError; "Expected int, got {}", safe_repr(value)); + } + } ObType::Float | ObType::FloatSubclass => { let v = value.extract::()?; if (v.is_nan() || v.is_infinite()) && extra.config.inf_nan_mode == InfNanMode::Null { diff --git a/tests/serializers/test_any.py b/tests/serializers/test_any.py index 12ac463e1..344e0ff02 100644 --- a/tests/serializers/test_any.py +++ b/tests/serializers/test_any.py @@ -654,6 +654,40 @@ def test_ser_json_inf_nan_with_list_of_any() -> None: assert s.to_json([nan]) == b'[null]' +def test_ser_json_int_subclass_value_larger_than_i64(): + class IntSubclass(int): + pass + + schema = core_schema.model_schema( + MyModel, + core_schema.model_fields_schema( + dict( + stuff=core_schema.model_field( + core_schema.dict_schema( + keys_schema=core_schema.str_schema(), + values_schema=core_schema.any_schema(), + ) + ) + ) + ), + ) + s = SchemaSerializer(schema) + + assert ( + s.to_json( + MyModel(stuff={'value': IntSubclass(9_223_372_036_854_775_809)}), + ) + == b'{"stuff":{"value":9223372036854775809}}' + ) + + assert str( + s.to_python( + MyModel(stuff={'value': IntSubclass(9_223_372_036_854_775_809)}), + mode='json', + ) + ) == str({'stuff': {'value': 9223372036854775809}}) + + def test_simple_any_ser_schema_repr(): assert ( plain_repr(SchemaSerializer(core_schema.simple_ser_schema('any'))) From 725d6404b9376c668a38698430bcf5b2ef37f263 Mon Sep 17 00:00:00 2001 From: Kangxing Xie Date: Sat, 24 Aug 2024 20:06:08 +1000 Subject: [PATCH 2/2] chore: re-added the white spaces to maintain the readability --- src/serializers/infer.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/serializers/infer.rs b/src/serializers/infer.rs index 18b21356e..00add47f3 100644 --- a/src/serializers/infer.rs +++ b/src/serializers/infer.rs @@ -45,6 +45,7 @@ pub(crate) fn infer_to_python_known( mut extra: &Extra, ) -> PyResult { let py = value.py(); + let mode = extra.mode; let mut guard = match extra.recursion_guard(value, INFER_DEF_REF_ID) { Ok(v) => v, @@ -110,6 +111,7 @@ pub(crate) fn infer_to_python_known( ); serializer.serializer.to_python(value, include, exclude, &extra) }; + let value = match extra.mode { SerMode::Json => match ob_type { // `bool` and `None` can't be subclasses, `ObType::Int`, `ObType::Float`, `ObType::Str` refer to exact types