fallback: Use mmap objects instead of strings to unpack

Signed-off-by: Bas Westerbaan <bas@westerbaan.name>
This commit is contained in:
Bas Westerbaan 2013-01-29 03:39:46 +01:00
parent b9e9199eea
commit 770fed6b7f

View file

@ -1,15 +1,27 @@
# Fallback pure Python implementation of msgpack # Fallback pure Python implementation of msgpack
#
# Easy imports
#
import sys import sys
import array import array
import struct import struct
#
# Tricky imports
#
try:
from cStringIO import StringIO
except ImportError:
from StringIO import StringIO
# We will use wStringIO for buffering the writes for packing.
# Normally, we will use cStringIO.StringIO.
# On PyPy we will use PyPy's own StringBuilder.
if hasattr(sys, 'pypy_version_info'): if hasattr(sys, 'pypy_version_info'):
# cStringIO is slow on PyPy, StringIO is faster. However: PyPy's own
# StringBuilder is fastest.
from __pypy__.builders import StringBuilder from __pypy__.builders import StringBuilder
USING_STRINGBUILDER = True USING_STRINGBUILDER = True
class StringIO(object): class wStringIO(object):
def __init__(self, s=''): def __init__(self, s=''):
if s: if s:
self.builder = StringBuilder(len(s)) self.builder = StringBuilder(len(s))
@ -22,10 +34,18 @@ if hasattr(sys, 'pypy_version_info'):
return self.builder.build() return self.builder.build()
else: else:
USING_STRINGBUILDER = False USING_STRINGBUILDER = False
try: wStringIO = StringIO
from cStringIO import StringIO
except ImportError: # We will use rStringIO for unpacking.
from StringIO import StringIO # Normally, this is a mmap. A normal StringIO is not a drop-in replacement ---
# it misses the __len__ operation.
# TODO add fallback for when mmap is unavailable
import mmap
def rStringIO(s):
m = mmap.mmap(-1, len(s))
m.write(s)
m.seek(0)
return m
from msgpack.exceptions import ( from msgpack.exceptions import (
BufferFull, BufferFull,
@ -184,13 +204,13 @@ class Unpacker(object):
if self._fb_buf_n + len(next_bytes) > self.max_buffer_size: if self._fb_buf_n + len(next_bytes) > self.max_buffer_size:
raise BufferFull raise BufferFull
self._fb_buf_n += len(next_bytes) self._fb_buf_n += len(next_bytes)
self._fb_buffers.append(next_bytes) self._fb_buffers.append(rStringIO(next_bytes))
def _fb_consume(self): def _fb_consume(self):
self._fb_buffers = self._fb_buffers[self._fb_buf_i:] self._fb_buffers = self._fb_buffers[self._fb_buf_i:]
if self._fb_buffers: if self._fb_buffers:
self._fb_buffers[0] = self._fb_buffers[0][self._fb_buf_o:] self._fb_buffers[0] = rStringIO(self._fb_buffers[0][
self._fb_buf_o = 0 self._fb_buffers[0].tell():])
self._fb_buf_i = 0 self._fb_buf_i = 0
self._fb_buf_n = sum(map(len, self._fb_buffers)) self._fb_buf_n = sum(map(len, self._fb_buffers))
@ -212,16 +232,20 @@ class Unpacker(object):
return self._fb_read(n) return self._fb_read(n)
def _fb_rollback(self): def _fb_rollback(self):
for buf in self._fb_buffers:
buf.seek(0)
self._fb_buf_i = 0 self._fb_buf_i = 0
self._fb_buf_o = 0
def _fb_get_extradata(self): def _fb_get_extradata(self):
bufs = self._fb_buffers[self._fb_buf_i:] bufs = self._fb_buffers[self._fb_buf_i:]
if bufs: if bufs:
bufs[0] = bufs[0][self._fb_buf_o:] bufs[0] = rStringIO(bufs[0][bufs[0].tell():])
return ''.join(bufs) return ''.join([buf[:] for buf in bufs])
def _fb_read(self, n, write_bytes=None): def _fb_read(self, n, write_bytes=None):
if (write_bytes is None and self._fb_buf_i < len(self._fb_buffers)
and self._fb_buffers[0].tell() + n < len(self._fb_buffers[0])):
return self._fb_buffers[0].read(n)
ret = '' ret = ''
while len(ret) != n: while len(ret) != n:
if self._fb_buf_i == len(self._fb_buffers): if self._fb_buf_i == len(self._fb_buffers):
@ -230,14 +254,12 @@ class Unpacker(object):
tmp = self.file_like.read(self.read_size) tmp = self.file_like.read(self.read_size)
if not tmp: if not tmp:
break break
self._fb_buffers.append(tmp) self._fb_buffers.append(rStringIO(tmp))
continue continue
sliced = n - len(ret) sliced = n - len(ret)
ret += self._fb_buffers[self._fb_buf_i][ ret += self._fb_buffers[self._fb_buf_i].read(sliced)
self._fb_buf_o:self._fb_buf_o + sliced] if (self._fb_buffers[self._fb_buf_i].tell()
self._fb_buf_o += sliced == len(self._fb_buffers[self._fb_buf_i])):
if self._fb_buf_o >= len(self._fb_buffers[self._fb_buf_i]):
self._fb_buf_o = 0
self._fb_buf_i += 1 self._fb_buf_i += 1
if len(ret) != n: if len(ret) != n:
self._fb_rollback() self._fb_rollback()
@ -394,7 +416,7 @@ class Packer(object):
self.autoreset = autoreset self.autoreset = autoreset
self.encoding = encoding self.encoding = encoding
self.unicode_errors = unicode_errors self.unicode_errors = unicode_errors
self.buffer = StringIO() self.buffer = wStringIO()
if default is not None: if default is not None:
if not callable(default): if not callable(default):
raise TypeError("default must be callable") raise TypeError("default must be callable")
@ -464,33 +486,33 @@ class Packer(object):
self._pack(obj) self._pack(obj)
ret = self.buffer.getvalue() ret = self.buffer.getvalue()
if self.autoreset: if self.autoreset:
self.buffer = StringIO() self.buffer = wStringIO()
elif USING_STRINGBUILDER: elif USING_STRINGBUILDER:
self.buffer = StringIO(ret) self.buffer = wStringIO(ret)
return ret return ret
def pack_map_pairs(self, pairs): def pack_map_pairs(self, pairs):
self._fb_pack_map_pairs(len(pairs), pairs) self._fb_pack_map_pairs(len(pairs), pairs)
ret = self.buffer.getvalue() ret = self.buffer.getvalue()
if self.autoreset: if self.autoreset:
self.buffer = StringIO() self.buffer = wStringIO()
elif USING_STRINGBUILDER: elif USING_STRINGBUILDER:
self.buffer = StringIO(ret) self.buffer = wStringIO(ret)
return ret return ret
def pack_array_header(self, n): def pack_array_header(self, n):
self._fb_pack_array_header(n) self._fb_pack_array_header(n)
ret = self.buffer.getvalue() ret = self.buffer.getvalue()
if self.autoreset: if self.autoreset:
self.buffer = StringIO() self.buffer = wStringIO()
elif USING_STRINGBUILDER: elif USING_STRINGBUILDER:
self.buffer = StringIO(ret) self.buffer = wStringIO(ret)
return ret return ret
def pack_map_header(self, n): def pack_map_header(self, n):
self._fb_pack_map_header(n) self._fb_pack_map_header(n)
ret = self.buffer.getvalue() ret = self.buffer.getvalue()
if self.autoreset: if self.autoreset:
self.buffer = StringIO() self.buffer = wStringIO()
elif USING_STRINGBUILDER: elif USING_STRINGBUILDER:
self.buffer = StringIO(ret) self.buffer = wStringIO(ret)
return ret return ret
def _fb_pack_array_header(self, n): def _fb_pack_array_header(self, n):
if n <= 0x0f: if n <= 0x0f:
@ -516,4 +538,4 @@ class Packer(object):
def bytes(self): def bytes(self):
return self.buffer.getvalue() return self.buffer.getvalue()
def reset(self): def reset(self):
self.buffer = StringIO() self.buffer = wStringIO()