__all__ = ['Session']
classSession(object):def__init__(self):
self.clients = {}
defget(self, client):# Get transport associated by client if exists.if client notin self.clients:
returnNonereturn self.clients[client]
def__contains__(self, client):# Decide if client is onlinereturn client in self.clients
def__repr__(self):return"{}".format(self.clients)
__str__ = __repr__
defregister(self, client, transport):"""Register client on session"""
self.clients[client] = transport
defunregister(self, client):"""Unregister client on session"""if client in self.clients:
del self.clients[client]
if __name__ == '__main__':
Session()
handlers.py
__all__ = ['MessageHandler']
import asyncio
import json
from struct import pack
from session import Session
classMetaHandler(type):"""Metaclass for MessageHandler"""def__init__(cls, name, bases, _dict):try:
cls._msg_handlers[cls.__msgtype__] = cls
except AttributeError:
cls._msg_handlers = {}
classMessageHandler(metaclass=MetaHandler):
_session = Session()
defhandle(self, msg, transport):try:
_handler = self._msg_handlers[msg['type']]
except KeyError:
return ErrorHandler().handler(msg)
# Handling messages in a asyncio-Task# Don’t directly create Task instances: use the async() function# or the BaseEventLoop.create_task() method.#return _handler().handle(msg, transport)return asyncio.async(_handler().handle(msg, transport))
classErrorHandler(MessageHandler):"""
Unknown message type
"""
__msgtype__ = 'unknown'defhandle(self, msg):
print("Unknown message type: {}".format(msg))
classRegister(MessageHandler):"""
Registry handler for handling clients registry.
Message body should like this:
{'type': 'register', 'uid': 'unique-user-id'}
"""
__msgtype__ = 'register'def__init__(self):
self.current_uid = None
self.transport = None@asyncio.coroutinedefhandle(self, msg, transport):
self.current_uid = msg['uid']
self.transport = transport
print("registe uid: {}".format(self.current_uid))
# Register user in global session
self._session.register(self.current_uid, self.transport)
classSendTextMsg(MessageHandler):"""
Send message to others.
Message body should like this:
{'type': 'text', 'sender': 'Jack', 'receiver': 'Rose', 'content': 'I love you forever'}
"""
__msgtype__ = 'text'# Text message@asyncio.coroutinedefhandle(self, msg, _):"""
Send message to receiver if receiver is online, and
save message to mongodb. Otherwise save
message to mongodb as offline message.
:param msg:
:return: None
"""
print("send data...{}".format(msg))
transport = self._session.get(msg['receiver'])
msg_pack = json.dumps(msg)
msg_len = len(msg_pack)
if transport:
# Pack message as length-prifixed and send to receiver.
transport.write(pack("!I%ds" % msg_len, msg_len, bytes(msg_pack, encoding='utf-8')))
classUnregister(MessageHandler):"""
Unregister user from global session
Message body should like this:
{'type': 'unregister', 'uid': 'unique-user-id'}
"""
__msgtype__ = 'unregister'@asyncio.coroutinedefhandle(self, msg, _):"""Unregister user record from global session"""
self._session.unregister(msg['uid'])
server.py
import asyncio
import json
from handlers import MessageHandler
# 消息頭長度 int or uint
_MESSAGE_PREFIX_LENGTH = 4# 字節序
_BYTE_ORDER = 'big'classmyImProtocol(asyncio.Protocol):
_buffer = b''# 數據緩衝Buffer
_msg_len = None# 消息長度defdata_received(self, data):while data:
data = self.process_data(data)
defprocess_data(self, data):"""
Called when some data is received.
This method must be implemented by subclasses
The argument is a bytes object.
"""
self._buffer += data
# For store the rest data out-of a full message
_buffer = Noneif self._msg_len isNone:
# If buffer length < _MESSAGE_PREFIX_LENGTH return for more dataif len(self._buffer) < _MESSAGE_PREFIX_LENGTH:
return# If buffer length >= _MESSAGE_PREFIX_LENGTH
self._msg_len = int.from_bytes(self._buffer[:_MESSAGE_PREFIX_LENGTH], byteorder=_BYTE_ORDER)
# The left bytes will be the message body
self._buffer = self._buffer[_MESSAGE_PREFIX_LENGTH:]
# Received full messageif len(self._buffer) >= self._msg_len:
# Call message_received to handler message
self.message_received(self._buffer[:self._msg_len])
# Left the rest of the buffer for next message
_buffer = self._buffer[self._msg_len:]
# Clean data buffer for next message
self._buffer = b''# Set message length to None for next message
self._msg_len = Nonereturn _buffer
defmessage_received(self, msg):"""
Must override in subclass
:param msg: the full message
:return: None
"""raise NotImplementedError()
classmyIm(myImProtocol):def__init__(self):
self.handler = MessageHandler()
self.transport = Nonedefconnection_made(self, transport):
self.transport = transport
defmessage_received(self, msg):"""
The real message handler
:param msg: a full message without prefix length
:return: None
"""# Convert bytes msg to python dictionary
msg = json.loads(msg.decode("utf-8"))
print("receive msg...{}".format(msg))
# Handler msgreturn self.handler.handle(msg, self.transport)
classmyImServer(object):def__init__(self, protocol_factory, host, port):
self.host = host
self.port = port
self.protocol_factory = protocol_factory
defstart(self):
loop = asyncio.get_event_loop()
loop.run_until_complete(loop.create_server(self.protocol_factory, self.host, self.port))
loop.run_forever()
if __name__ == '__main__':
server = myImServer(myIm, 'localhost', 2222)
server.start()