From e37436287dd29b826a5077042fb2eec7a24d2be3 Mon Sep 17 00:00:00 2001 From: Peter Chang Date: Mon, 20 Jul 2026 17:44:29 +0100 Subject: [PATCH 1/2] Allow default handler to take a position parameter This enables padding to be computed when packing a custom type and enables zero-copy unpacking by aligning primitives correctly --- msgpack/_packer.pyx | 18 ++++++- msgpack/fallback.py | 28 ++++++++--- test/test_default.py | 114 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 152 insertions(+), 8 deletions(-) create mode 100644 test/test_default.py diff --git a/msgpack/_packer.pyx b/msgpack/_packer.pyx index 277239d8..380b336b 100644 --- a/msgpack/_packer.pyx +++ b/msgpack/_packer.pyx @@ -8,6 +8,7 @@ from cpython.datetime cimport ( cdef ExtType cdef Timestamp +import inspect from .ext import ExtType, Timestamp @@ -65,7 +66,8 @@ cdef class Packer: :param default: When specified, it should be callable. Convert user type to builtin type that Packer supports. - See also simplejson's document. + See also simplejson's document. In addition, this callable may have two parameters + where the second parameter is given the current position of the stream. :param bool use_single_float: Use single precision float type for float. (default: False) @@ -106,6 +108,7 @@ cdef class Packer: cdef const char *unicode_errors cdef size_t exports # number of exported buffers cdef bint strict_types + cdef bint _pass_posn cdef bint use_float cdef bint autoreset cdef bint datetime @@ -140,6 +143,14 @@ cdef class Packer: if default is not None: if not PyCallable_Check(default): raise TypeError("default must be a callable.") + default_argc = len(inspect.signature(default).parameters) + if default_argc == 1: + self._pass_posn = False + elif default_argc == 2: + self._pass_posn = True + else: + raise ValueError("default must take one or two parameters") + self._default = default self._berrors = unicode_errors @@ -263,7 +274,10 @@ cdef class Packer: if self._default is not None: ret = self._pack_inner(o, 1, nest_limit) if ret == -2: - o = self._default(o) + if self._pass_posn: + o = self._default(o, self.pk.length) + else: + o = self._default(o) else: return ret return self._pack_inner(o, 0, nest_limit) diff --git a/msgpack/fallback.py b/msgpack/fallback.py index 824f59d5..08b90ce7 100644 --- a/msgpack/fallback.py +++ b/msgpack/fallback.py @@ -1,5 +1,6 @@ """Fallback pure Python implementation of msgpack""" +import inspect import struct import sys from datetime import datetime as _DateTime @@ -622,7 +623,8 @@ class Packer: :param default: When specified, it should be callable. Convert user type to builtin type that Packer supports. - See also simplejson's document. + See also simplejson's document. In addition, this callable may have two parameters + where the second parameter is given the current position of the stream. :param bool use_single_float: Use single precision float type for float. (default: False) @@ -675,8 +677,16 @@ def __init__( self._buffer = BytesIO() self._datetime = bool(datetime) self._unicode_errors = unicode_errors or "strict" - if default is not None and not callable(default): - raise TypeError("default must be callable") + if default is not None: + if not callable(default): + raise TypeError("default must be callable") + default_argc = len(inspect.signature(default).parameters) + if default_argc == 1: + self._pass_posn = False + elif default_argc == 2: + self._pass_posn = True + else: + raise ValueError("default must take one or two parameters") self._default = default def _pack( @@ -723,8 +733,11 @@ def _pack( if -0x8000000000000000 <= obj < -0x80000000: return self._buffer.write(struct.pack(">Bq", 0xD3, obj)) if not default_used and self._default is not None: - obj = self._default(obj) default_used = True + if self._pass_posn: + obj = self._default(obj, self._buffer.tell()) + else: + obj = self._default(obj) continue raise OverflowError("Integer value out of range") if check(obj, (bytes, bytearray)): @@ -794,8 +807,11 @@ def _pack( continue if not default_used and self._default is not None: - obj = self._default(obj) - default_used = 1 + default_used = True + if self._pass_posn: + obj = self._default(obj, self._buffer.tell()) + else: + obj = self._default(obj) continue if self._datetime and check(obj, _DateTime): diff --git a/test/test_default.py b/test/test_default.py new file mode 100644 index 00000000..b7d430f8 --- /dev/null +++ b/test/test_default.py @@ -0,0 +1,114 @@ +import logging +from array import array +from typing import Any + +import msgpack + + +def dfa(o: array): + logging.info("Default called for %s", o) + return dict(data=o.tobytes(), tcode=o.typecode) + + +def decode_array(obj): + if isinstance(obj, dict) and "tcode" in obj: + return memoryview(obj["data"]).cast(obj["tcode"]) + return obj + + +def encode_array(obj: Any, posn: int) -> Any | dict[str, Any]: + if isinstance(obj, array): + tcode = obj.typecode + itemsize = obj.itemsize + length = len(obj) + bytesize = itemsize * length + offset = posn + 1 + 5 + 1 # fixmap, fixstr + 4, binX + if bytesize < 256: + offset += 1 + elif bytesize < 65536: + offset += 2 + elif bytesize < 2**32: + offset += 4 + else: + raise ValueError("Array too large (>= 4GiB)") + + print(f"{offset=}, {itemsize=}") + if offset % itemsize == 0: + pad = 0 + else: + offset += 4 # prefix with pad so fixstr + 3 + ? + pad = itemsize - (offset % itemsize) + # use fixstr for padding so deduct one and wrap + pad = (pad - 1) % itemsize + + if pad: + obj = dict( + pad="." * pad, + data=obj.tobytes(), + tcode=tcode, + ) + else: + obj = dict( + data=obj.tobytes(), + tcode=tcode, + ) + return obj + + +def check_align(pa, itemsize): + """ + check of start of data, 0xc4 + 1/0xc5 + 2, 0xc6 + 4 + has aligned offset + """ + offset = pa.index(b"\xa4data") + 5 + print(f"After data={offset}") + if pa.startswith(b"\xc4", offset): + offset += 2 + elif pa.startswith(b"\xc5", offset): + offset += 3 + elif pa.startswith(b"\xc6", offset): + offset += 5 + print(f"{offset=}, align={offset % itemsize}") + assert offset % itemsize == 0 + + +def test_default_top(): + eo = array("f", [0.1]) + pa = msgpack.packb(eo, default=dfa) + print(pa) + ao = msgpack.unpackb(pa, object_hook=decode_array) + assert eo == ao + + pa = msgpack.packb(eo, default=encode_array) + print(pa) + check_align(pa, eo.itemsize) + ad = msgpack.unpackb(pa) + assert eo == memoryview(ad["data"]).cast(ad["tcode"]) + ao = msgpack.unpackb(pa, object_hook=decode_array) + assert eo == ao + print("Unpacked as", ao.tolist()) + + +def test_default_second(): + eo = dict(config="ladybug", ndata=array("d", [0.1])) + pa = msgpack.packb(eo, default=dfa) + print(pa) + ao = msgpack.unpackb(pa, object_hook=decode_array) + assert eo == ao + + pa = msgpack.packb(eo, default=encode_array) + print(pa) + check_align(pa, eo["ndata"].itemsize) + print(eo) + ao = msgpack.unpackb(pa) + ad = ao["ndata"] + assert eo["ndata"] == memoryview(ad["data"]).cast(ad["tcode"]) + ao = msgpack.unpackb(pa, object_hook=decode_array) + assert eo == ao + print("Unpacked as", ao, ao["ndata"].tolist()) + + +if __name__ == "__main__": + print(msgpack.__file__) + test_default_top() + test_default_second() From 1fd4ac4ce4ce1c7dbb83dcd363b32381f309e220 Mon Sep 17 00:00:00 2001 From: Peter Chang Date: Wed, 22 Jul 2026 13:56:56 +0100 Subject: [PATCH 2/2] Fix pad wrapping in test --- test/test_default.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/test_default.py b/test/test_default.py index b7d430f8..3db367cc 100644 --- a/test/test_default.py +++ b/test/test_default.py @@ -38,8 +38,10 @@ def encode_array(obj: Any, posn: int) -> Any | dict[str, Any]: else: offset += 4 # prefix with pad so fixstr + 3 + ? pad = itemsize - (offset % itemsize) - # use fixstr for padding so deduct one and wrap - pad = (pad - 1) % itemsize + # use fixstr for padding so deduct one + if pad < 2: + pad += itemsize # prewrap as pad cannot be less than one + pad -= 1 if pad: obj = dict(