diff --git a/src/socketio/msgpack_packet.py b/src/socketio/msgpack_packet.py index 27462634..32c4debd 100644 --- a/src/socketio/msgpack_packet.py +++ b/src/socketio/msgpack_packet.py @@ -1,18 +1,56 @@ +import logging import msgpack from . import packet +logger = logging.getLogger('socketio') + class MsgPackPacket(packet.Packet): uses_binary_events = False def encode(self): """Encode the packet for transmission.""" - return msgpack.dumps(self._to_dict()) + return self._encode() + + def _encode(self, **kwargs): + return _msgpack.dumps(self._to_dict(), **kwargs) def decode(self, encoded_packet): """Decode a transmitted package.""" - decoded = msgpack.loads(encoded_packet) + return self._decode(encoded_packet) + + def _decode(self, encoded_packet, **kwargs): + decoded = msgpack.loads(encoded_packet, **kwargs) self.packet_type = decoded['type'] self.data = decoded.get('data') self.id = decoded.get('id') self.namespace = decoded['nsp'] + + @classmethod + def _configure(cls, *args, **kwargs): + dumps_default = kwargs.pop('dumps_default', None) + ext_hook = kwargs.pop('ext_hook', msgpack.ExtType) + + if args: + logger.warning( + 'Some positional arguments to MsgPackPacket.configure() are ' + 'not used: %s', + args, + ) + if kwargs: + logger.warning( + 'Some keyword arguments to MsgPackPacket.configure() are ' + 'not used: %s', + kwargs, + ) + + class ConfiguredMsgPackPacket(cls): + def _encode(self, **kwargs): + kwargs.setdefault('default', dumps_default) + return super()._encode(**kwargs) + + def _decode(self, encoded_packet, **kwargs): + kwargs.setdefault('ext_hook', ext_hook) + return super()._decode(encoded_packet, **kwargs) + + return ConfiguredMsgPackPacket diff --git a/src/socketio/packet.py b/src/socketio/packet.py index 3deba7fb..421c800b 100644 --- a/src/socketio/packet.py +++ b/src/socketio/packet.py @@ -1,4 +1,5 @@ import functools +import logging from engineio import json as _json (CONNECT, DISCONNECT, EVENT, ACK, CONNECT_ERROR, BINARY_EVENT, BINARY_ACK) = \ @@ -6,6 +7,7 @@ packet_names = ['CONNECT', 'DISCONNECT', 'EVENT', 'ACK', 'CONNECT_ERROR', 'BINARY_EVENT', 'BINARY_ACK'] +logger = logging.getLogger('socketio.packet') class Packet: """Socket.IO packet.""" @@ -21,6 +23,8 @@ class Packet: uses_binary_events = True json = _json + _configure_args = ((),()) + _subclass_registry = {} def __init__(self, packet_type=EVENT, data=None, namespace=None, id=None, binary=None, encoded_packet=None): @@ -192,3 +196,29 @@ def _to_dict(self): if self.id is not None: d['id'] = self.id return d + + @classmethod + def configure(cls, *args, **kwargs): + configure_args = (args, tuple(sorted(kwargs.items()))) + try: + args_hash = hash(configure_args) + except TypeError: + logger.warning("Packet.configure() called with unhashable " + "arguments; subclass caching will not work.") + args_hash = None + + if args_hash in cls._subclass_registry: + logger.debug("Using cached Packet subclass for args %s, %s", + args, kwargs) + return cls._subclass_registry[args_hash] + new = cls._configure(*args, **kwargs) + if args_hash is not None: + cls._subclass_registry[args_hash] = new + logger.debug("Caching Packet subclass for args %s, %s", + args, kwargs) + return new + + @classmethod + def _configure(cls, *args, **kwargs): + raise NotImplementedError('Packet._configure() must be implemented ' + 'by subclasses.') \ No newline at end of file