refactor / optimize

Author: gesslerpd

scripts/bench.py

@@ -0,0 +1,61 @@
+import os
+
+# set before importing `msgpack`
+os.environ["MSGPACK_PUREPYTHON"] = "1"
+
+from msgpack_stream import unpack, unpack_stream
+from mmap import mmap, ACCESS_READ
+import argparse
+import timeit
+
+from msgpack import unpackb
+
+
+FILE = "scripts/obj.msgpack"
+
+
+def main(mapped):
+    with open(FILE, "rb", buffering=0) as fd:
+        if mapped:
+            with mmap(fd.fileno(), 0, access=ACCESS_READ) as mm:
+                return unpack(mm.read())
+        else:
+            return unpack(fd.read())
+
+
+def stream(mapped):
+    with open(FILE, "rb", buffering=0) as fd:
+        if mapped:
+            with mmap(fd.fileno(), 0, access=ACCESS_READ) as mm:
+                return unpack_stream(mm)
+        else:
+            return unpack_stream(fd)
+
+
+def other(mapped):
+    with open(FILE, "rb", buffering=0) as fd:
+        if mapped:
+            with mmap(fd.fileno(), 0, access=ACCESS_READ) as mm:
+                return unpackb(mm.read(), strict_map_key=False)
+        else:
+            return unpackb(fd.read(), strict_map_key=False)
+
+
+if __name__ == "__main__":
+    _globals = {
+        "main": main,
+        "stream": stream,
+        "other": other,
+    }
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--number", type=int, default=25, help="Number of runs")
+    args = parser.parse_args()
+
+    t_main = timeit.timeit("main(True)", number=args.number, globals=_globals)
+    # this needs to be mmap for good performance
+    t_stream = timeit.timeit("stream(True)", number=args.number, globals=_globals)
+    t_other = timeit.timeit("other(True)", number=args.number, globals=_globals)
+
+    print(f"main: {t_main:.6f}s total, {t_main / args.number:.6f}s per call")
+    print(f"stream: {t_stream:.6f}s total, {t_stream / args.number:.6f}s per call")
+    print(f"other: {t_other:.6f}s total, {t_other / args.number:.6f}s per call")

scripts/obj.msgpack

@@ -1,14 +1,4 @@
 """Pure Python stream based implementation of msgpack"""
 
-from ._io import pack as packb, unpack as unpackb
-from ._msgpack import msgpack_t
-
-
-def pack(obj):
-    """Pack object into data."""
-    return packb(msgpack_t, obj)
-
-
-def unpack(data):
-    """Unpack data into object."""
-    return unpackb(msgpack_t, data)
+from ._io import pack as pack, unpack as unpack
+from ._msgpack import pack_stream as pack_stream, unpack_stream as unpack_stream

src/msgpack_stream/__init__.py

@@ -1,25 +1,19 @@
 import io
 
+from ._msgpack import pack_stream, unpack_stream
 
-class container(dict):
-    __getattr__ = dict.__getitem__
-    __setattr__ = dict.__setitem__
-    __delattr__ = dict.__delitem__
 
-
-def pack(typ, obj):
-    stream = io.BytesIO()
-    typ.pack(stream, obj)
-    data = stream.getvalue()
-    stream.close()
+def pack(obj):
+    """Pack object into data."""
+    with io.BytesIO() as stream:
+        pack_stream(stream, obj)
+        data = stream.getvalue()
     return data
 
 
-def unpack(typ, data):
-    stream = io.BytesIO(data)
-    obj = typ.unpack(stream)
-    extra_data = stream.read()
-    if extra_data:
-        raise RuntimeError("too much data", extra_data)
-    stream.close()
-    return obj
+def unpack(data):
+    """Unpack data into object."""
+    with io.BytesIO(data) as stream:
+        obj = unpack_stream(stream)
+        excess_data = stream.read()
+    return obj, excess_data