Allow default handler to take a position parameter

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)

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):

test/test_default.py

@@ -0,0 +1,116 @@
+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
+            if pad < 2:
+                pad += itemsize # prewrap as pad cannot be less than one
+            pad -= 1
+
+        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()