packer: add buf_size option (#604)

methane · rtobar · web-flow · commit 72e65feb0e02 · 2024-05-06T00:49:12.000+09:00
And change the default buffer size to 256KiB.

Signed-off-by: Rodrigo Tobar &lt;rtobar@icrar.org&gt;
Co-authored-by: Rodrigo Tobar &lt;rtobar@icrar.org&gt;
diff --git a/msgpack/_packer.pyx b/msgpack/_packer.pyx
@@ -53,7 +53,7 @@ cdef inline int PyBytesLike_CheckExact(object o):
     return PyBytes_CheckExact(o) or PyByteArray_CheckExact(o)
 
 
-cdef class Packer(object):
+cdef class Packer:
     """
     MessagePack Packer
 
@@ -97,6 +97,11 @@ cdef class Packer(object):
     :param str unicode_errors:
         The error handler for encoding unicode. (default: 'strict')
         DO NOT USE THIS!!  This option is kept for very specific usage.
+
+    :param int buf_size:
+        The size of the internal buffer. (default: 256*1024)
+        Useful if serialisation size can be correctly estimated,
+        avoid unnecessary reallocations.
     """
     cdef msgpack_packer pk
     cdef object _default
@@ -107,8 +112,7 @@ cdef class Packer(object):
     cdef bint autoreset
     cdef bint datetime
 
-    def __cinit__(self):
-        cdef int buf_size = 1024*1024
+    def __cinit__(self, buf_size=256*1024, **_kwargs):
         self.pk.buf = <char*> PyMem_Malloc(buf_size)
         if self.pk.buf == NULL:
             raise MemoryError("Unable to allocate internal buffer.")
@@ -117,7 +121,8 @@ cdef class Packer(object):
 
     def __init__(self, *, default=None,
                  bint use_single_float=False, bint autoreset=True, bint use_bin_type=True,
-                 bint strict_types=False, bint datetime=False, unicode_errors=None):
+                 bint strict_types=False, bint datetime=False, unicode_errors=None,
+                 buf_size=256*1024):
         self.use_float = use_single_float
         self.strict_types = strict_types
         self.autoreset = autoreset
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
@@ -210,7 +210,7 @@ def unpackb(object packed, *, object object_hook=None, object list_hook=None,
     raise ValueError("Unpack failed: error = %d" % (ret,))
 
 
-cdef class Unpacker(object):
+cdef class Unpacker:
     """Streaming unpacker.
 
     Arguments:
diff --git a/msgpack/fallback.py b/msgpack/fallback.py
@@ -232,6 +232,7 @@ class Unpacker:
     def __init__(
         self,
         file_like=None,
+        *,
         read_size=0,
         use_list=True,
         raw=False,
@@ -650,39 +651,21 @@ class Packer:
         The error handler for encoding unicode. (default: 'strict')
         DO NOT USE THIS!!  This option is kept for very specific usage.
 
-    Example of streaming deserialize from file-like object::
-
-        unpacker = Unpacker(file_like)
-        for o in unpacker:
-            process(o)
-
-    Example of streaming deserialize from socket::
-
-        unpacker = Unpacker()
-        while True:
-            buf = sock.recv(1024**2)
-            if not buf:
-                break
-            unpacker.feed(buf)
-            for o in unpacker:
-                process(o)
-
-    Raises ``ExtraData`` when *packed* contains extra bytes.
-    Raises ``OutOfData`` when *packed* is incomplete.
-    Raises ``FormatError`` when *packed* is not valid msgpack.
-    Raises ``StackError`` when *packed* contains too nested.
-    Other exceptions can be raised during unpacking.
+    :param int buf_size:
+        Internal buffer size. This option is used only for C implementation.
     """
 
     def __init__(
         self,
+        *,
         default=None,
         use_single_float=False,
         autoreset=True,
         use_bin_type=True,
         strict_types=False,
         datetime=False,
         unicode_errors=None,
+        buf_size=None,
     ):
         self._strict_types = strict_types
         self._use_float = use_single_float
