From c17aaff9f9799032e72270870225e12d8a2f9a1f Mon Sep 17 00:00:00 2001 From: Tariq Bontekoe Date: Sat, 8 Jan 2022 11:33:23 +0100 Subject: [PATCH] optional passthrough of big ints and tuples --- CHANGELOG.md | 3 +++ ormsgpack.pyi | 2 ++ src/lib.rs | 12 +++++++++++ src/opt.rs | 11 ++++++++-- src/serialize/serializer.rs | 7 +++++-- tests/test_api.py | 4 ++-- tests/test_type.py | 40 ++++++++++++++++++++++++++++++++++++- 7 files changed, 72 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9a6e1abd..13babc64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Changelog ## Next Version +### Added +- Add optional passthrough for tuples. by [@TariqTNO](https://github.com/aviramha/ormsgpack/pull/64) +- Add optional passthrough for ints, that do not fit into an i64. by [@TariqTNO](https://github.com/aviramha/ormsgpack/pull/64) ### Changed - `opt` parameter can be `None`. ### Misc diff --git a/ormsgpack.pyi b/ormsgpack.pyi index ef1207a0..cb0a7fe1 100644 --- a/ormsgpack.pyi +++ b/ormsgpack.pyi @@ -16,9 +16,11 @@ class MsgpackEncodeError(TypeError): ... OPT_NAIVE_UTC: int OPT_OMIT_MICROSECONDS: int +OPT_PASSTHROUGH_BIG_INT: int OPT_PASSTHROUGH_DATACLASS: int OPT_PASSTHROUGH_DATETIME: int OPT_PASSTHROUGH_SUBCLASS: int +OPT_PASSTHROUGH_TUPLE: int OPT_SERIALIZE_NUMPY: int OPT_SERIALIZE_PYDANTIC: int OPT_NON_STR_KEYS: int diff --git a/src/lib.rs b/src/lib.rs index add881d4..eaab7dfa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -166,6 +166,12 @@ pub unsafe extern "C" fn PyInit_ormsgpack() -> *mut PyObject { "OPT_OMIT_MICROSECONDS\0", opt::OMIT_MICROSECONDS ); + opt!( + exported_objects, + mptr, + "OPT_PASSTHROUGH_BIG_INT\0", + opt::PASSTHROUGH_BIG_INT + ); opt!( exported_objects, mptr, @@ -196,6 +202,12 @@ pub unsafe extern "C" fn PyInit_ormsgpack() -> *mut PyObject { "OPT_SERIALIZE_PYDANTIC\0", opt::SERIALIZE_PYDANTIC ); + opt!( + exported_objects, + mptr, + "OPT_PASSTHROUGH_TUPLE\0", + opt::PASSTHROUGH_TUPLE + ); opt!(exported_objects, mptr, "OPT_UTC_Z\0", opt::UTC_Z); typeref::init_typerefs(); diff --git a/src/opt.rs b/src/opt.rs index 5c967c17..4284f443 100644 --- a/src/opt.rs +++ b/src/opt.rs @@ -13,18 +13,25 @@ pub const PASSTHROUGH_DATETIME: Opt = 1 << 7; pub const APPEND_NEWLINE: Opt = 1 << 8; pub const PASSTHROUGH_DATACLASS: Opt = 1 << 9; pub const SERIALIZE_PYDANTIC: Opt = 1 << 10; +pub const PASSTHROUGH_BIG_INT: Opt = 1 << 11; +pub const PASSTHROUGH_TUPLE: Opt = 1 << 12; -pub const NOT_PASSTHROUGH: Opt = - !(PASSTHROUGH_DATETIME | PASSTHROUGH_DATACLASS | PASSTHROUGH_SUBCLASS); +pub const NOT_PASSTHROUGH: Opt = !(PASSTHROUGH_BIG_INT + | PASSTHROUGH_DATACLASS + | PASSTHROUGH_DATETIME + | PASSTHROUGH_SUBCLASS + | PASSTHROUGH_TUPLE); pub const MAX_PACKB_OPT: i32 = (APPEND_NEWLINE | INDENT_2 | NAIVE_UTC | NON_STR_KEYS | OMIT_MICROSECONDS + | PASSTHROUGH_BIG_INT | PASSTHROUGH_DATETIME | PASSTHROUGH_DATACLASS | PASSTHROUGH_SUBCLASS + | PASSTHROUGH_TUPLE | SERIALIZE_NUMPY | SERIALIZE_PYDANTIC | UTC_Z) as i32; diff --git a/src/serialize/serializer.rs b/src/serialize/serializer.rs index 5b1a3f23..61415dc5 100644 --- a/src/serialize/serializer.rs +++ b/src/serialize/serializer.rs @@ -70,7 +70,9 @@ pub fn pyobject_to_obtype(obj: *mut pyo3::ffi::PyObject, opts: Opt) -> ObType { ObType::Str } else if ob_type == BYTES_TYPE { ObType::Bytes - } else if ob_type == INT_TYPE { + } else if ob_type == INT_TYPE + && (opts & PASSTHROUGH_BIG_INT == 0 || ffi!(_PyLong_NumBits(obj)) <= 63) + { ObType::Int } else if ob_type == BOOL_TYPE { ObType::Bool @@ -104,7 +106,7 @@ pub fn pyobject_to_obtype_unlikely(obj: *mut pyo3::ffi::PyObject, opts: Opt) -> ObType::Date } else if ob_type == TIME_TYPE && opts & PASSTHROUGH_DATETIME == 0 { ObType::Time - } else if ob_type == TUPLE_TYPE { + } else if ob_type == TUPLE_TYPE && opts & PASSTHROUGH_TUPLE == 0 { ObType::Tuple } else if ob_type == UUID_TYPE { ObType::Uuid @@ -116,6 +118,7 @@ pub fn pyobject_to_obtype_unlikely(obj: *mut pyo3::ffi::PyObject, opts: Opt) -> ObType::StrSubclass } else if opts & PASSTHROUGH_SUBCLASS == 0 && is_subclass!(ob_type, Py_TPFLAGS_LONG_SUBCLASS) + && (opts & PASSTHROUGH_BIG_INT == 0 || ffi!(_PyLong_NumBits(obj)) <= 63) { ObType::Int } else if opts & PASSTHROUGH_SUBCLASS == 0 diff --git a/tests/test_api.py b/tests/test_api.py index ddc9a4be..0b8aab35 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -103,9 +103,9 @@ def test_option_range_high(): packb/unpackb() option out of range high """ with pytest.raises(ormsgpack.MsgpackEncodeError): - ormsgpack.packb(True, option=1 << 12) + ormsgpack.packb(True, option=1 << 14) with pytest.raises(ormsgpack.MsgpackDecodeError): - ormsgpack.unpackb("\x00", option=1 << 12) + ormsgpack.unpackb("\x00", option=1 << 14) def test_opts_multiple(): diff --git a/tests/test_type.py b/tests/test_type.py index 04016ad4..590bcf02 100644 --- a/tests/test_type.py +++ b/tests/test_type.py @@ -174,7 +174,7 @@ def test_int_64(value): assert ormsgpack.unpackb(ormsgpack.packb(value)) == value -@pytest.mark.parametrize("value", (9223372036854775807, -9223372036854775807)) +@pytest.mark.parametrize("value", (9223372036854775808, 18446744073709551615)) def test_uint_64(value): """ uint 64-bit @@ -190,6 +190,34 @@ def test_int_128(value): pytest.raises(ormsgpack.MsgpackEncodeError, ormsgpack.packb, value) +@pytest.mark.parametrize("value", (9223372036854775807, -9223372036854775807)) +def test_int_64_passthrough(value): + """ + int 64-bit with passthrough + """ + assert ormsgpack.unpackb(ormsgpack.packb(value, option=ormsgpack.OPT_PASSTHROUGH_BIG_INT)) == value + + +@pytest.mark.parametrize("value", (9223372036854775808, 18446744073709551615)) +def test_uint_64_passthrough(value): + """ + uint 64-bit with passthrough + """ + result = ormsgpack.unpackb(ormsgpack.packb(value, option=ormsgpack.OPT_PASSTHROUGH_BIG_INT, default=lambda x: {"int": x.to_bytes(16, "little", signed=True)})) + assert list(result.keys()) == ["int"] + assert int.from_bytes(result["int"], "little", signed=True) == value + + +@pytest.mark.parametrize("value", (18446744073709551616, -9223372036854775809)) +def test_int_128_passthrough(value): + """ + int 128-bit with passthrough + """ + result = ormsgpack.unpackb(ormsgpack.packb(value, option=ormsgpack.OPT_PASSTHROUGH_BIG_INT, default=lambda x: {"int": x.to_bytes(16, "little", signed=True)})) + assert list(result.keys()) == ["int"] + assert int.from_bytes(result["int"], "little", signed=True) == value + + @pytest.mark.parametrize( "value", ( @@ -282,6 +310,16 @@ def test_tuple(): assert ormsgpack.unpackb(ormsgpack.packb(obj)) == list(obj) +def test_tuple_passthrough(): + """ + tuple with passthrough + """ + obj = ("a", "😊", True, {"b": 1.1}, 2) + result = ormsgpack.unpackb(ormsgpack.packb(obj, option=ormsgpack.OPT_PASSTHROUGH_TUPLE, default=lambda x: {"tuple": list(x)})) + assert list(result.keys()) == ["tuple"] + assert tuple(result["tuple"]) == obj + + def test_dict(): """ dict