Allow default handler to take a position parameter by PeterC-DLS · Pull Request #718 · msgpack/msgpack-python · GitHub
Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions msgpack/__init__.py
18 changes: 16 additions & 2 deletions msgpack/_packer.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ from cpython.datetime cimport (
cdef ExtType
cdef Timestamp

import inspect
from .ext import ExtType, Timestamp


Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
28 changes: 22 additions & 6 deletions msgpack/fallback.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Fallback pure Python implementation of msgpack"""

import inspect
import struct
import sys
from datetime import datetime as _DateTime
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)):
Expand Down Expand Up @@ -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):
Expand Down
116 changes: 116 additions & 0 deletions test/test_default.py