Support BigInt serialization for int subclasses (#1417)

pydantic · Aug 26, 2024 · d0384c7 · d0384c7
1 parent f4a0675
commit d0384c7
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 4 deletions.
diff --git 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::*;
@@ -116,10 +117,15 @@ pub(crate) fn infer_to_python_known(
             // `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::<BigInt>() {
+                    b.into_py(py)
+                } else {
+                    return py_err!(PyTypeError; "Expected int, got {}", safe_repr(value));
+                }
+            }
             ObType::Float | ObType::FloatSubclass => {
                 let v = value.extract::<f64>()?;
                 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
@@ -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')))