From 9dac272f0be89998b6d73fd8fa7320b7a5250160 Mon Sep 17 00:00:00 2001 From: Andrei Date: Wed, 21 Sep 2016 22:37:32 -0500 Subject: [PATCH] use skema, etc --- disco/gateway/client.py | 67 +++++++++++++++++++++++++++------------- disco/gateway/events.py | 22 +++++++++++++ disco/gateway/packets.py | 16 +--------- disco/types/__init__.py | 0 disco/types/guild.py | 7 +++++ disco/types/user.py | 5 +++ requirements.txt | 25 +++++++++++++++ 7 files changed, 106 insertions(+), 36 deletions(-) create mode 100644 disco/gateway/events.py create mode 100644 disco/types/__init__.py create mode 100644 disco/types/guild.py create mode 100644 disco/types/user.py create mode 100644 requirements.txt diff --git a/disco/gateway/client.py b/disco/gateway/client.py index d408c57..7398033 100644 --- a/disco/gateway/client.py +++ b/disco/gateway/client.py @@ -1,10 +1,13 @@ import websocket import gevent import json +import zlib -from disco.gateway.packets import ( - Packet, DispatchPacket, HeartbeatPacket, ReconnectPacket, InvalidSessionPacket, HelloPacket, HeartbeatAckPacket, - ResumePacket, IdentifyPacket) +from holster.emitter import Emitter +# from holster.util import SimpleObject + +from disco.gateway.packets import OPCode, HeartbeatPacket, ResumePacket, IdentifyPacket +from disco.gateway.events import GatewayEvent from disco.util.logging import LoggingClass GATEWAY_VERSION = 6 @@ -24,6 +27,7 @@ class GatewayClient(LoggingClass): def __init__(self, client): super(GatewayClient, self).__init__() self.client = client + self.emitter = Emitter(gevent.spawn) # Websocket connection self.ws = None @@ -51,9 +55,25 @@ class GatewayClient(LoggingClass): self.send(HeartbeatPacket(data=self.seq)) gevent.sleep(interval / 1000) + def handle_dispatch(self, packet): + obj = GatewayEvent.from_dispatch(packet) + self.log.info('Got dispatch for %s', obj.user.id) + + def handle_heartbeat(self, packet): + pass + + def handle_reconnect(self, packet): + pass + + def handle_invalid_session(self, packet): + pass + def handle_hello(self, packet): self.log.info('Recieved HELLO, starting heartbeater...') - self._heartbeat_task = gevent.spawn(self.heartbeat_task, packet.heartbeat_interval) + self._heartbeat_task = gevent.spawn(self.heartbeat_task, packet['d']['heartbeat_interval']) + + def handle_heartbeat_ack(self, packet): + pass def connect(self): if not self._cached_gateway_url: @@ -69,30 +89,35 @@ class GatewayClient(LoggingClass): ) def on_message(self, ws, msg): - # TODO: ZLIB + # Check if we're JSON + if msg[0] != '{': + print 'zlib' + msg = zlib.decompress(msg) try: - packet = Packet.load_json(json.loads(msg)) - if packet.seq and packet.seq > self.seq: - self.seq = packet.seq + data = json.loads(msg) except: self.log.exception('Failed to load dispatch:') return - if isinstance(packet, DispatchPacket): - self.handle_dispatch(packet) - elif isinstance(packet, HeartbeatPacket): - self.handle_heartbeat(packet) - elif isinstance(packet, ReconnectPacket): - self.handle_reconnect(packet) - elif isinstance(packet, InvalidSessionPacket): - self.handle_invalid_session(packet) - elif isinstance(packet, HelloPacket): - self.handle_hello(packet) - elif isinstance(packet, HeartbeatAckPacket): - self.handle_heartbeat_ack(packet) + # Update sequence + if data['s'] and data['s'] > self.seq: + self.seq = data['s'] + + if data['op'] == OPCode.DISPATCH: + self.handle_dispatch(data) + elif data['op'] == OPCode.HEARTBEAT: + self.handle_heartbeat(data) + elif data['op'] == OPCode.RECONNECT: + self.handle_reconnect(data) + elif data['op'] == OPCode.INVALID_SESSION: + self.handle_invalid_session(data) + elif data['op'] == OPCode.HELLO: + self.handle_hello(data) + elif data['op'] == OPCode.HEARTBEAT_ACK: + self.handle_heartbeat_ack(data) else: - raise Exception('Unknown packet: {}'.format(packet)) + raise Exception('Unknown packet: {}'.format(data['op'])) def on_error(self, ws, error): print 'error', error diff --git a/disco/gateway/events.py b/disco/gateway/events.py new file mode 100644 index 0000000..6a4ce71 --- /dev/null +++ b/disco/gateway/events.py @@ -0,0 +1,22 @@ +import inflection +import skema + +from disco.types.user import User +from disco.types.guild import Guild + + +class GatewayEvent(skema.Model): + @staticmethod + def from_dispatch(obj): + cls = globals().get(inflection.camelize(obj['t'].lower())) + if not cls: + raise Exception('Could not find cls for {}'.format(obj['t'])) + + return cls(obj['d']) + + +class Ready(GatewayEvent): + version = skema.IntType(stored_name='v') + session_id = skema.StringType() + user = skema.ModelType(User) + guilds = skema.ListType(skema.ModelType(Guild)) diff --git a/disco/gateway/packets.py b/disco/gateway/packets.py index b4dfdd0..8a8e818 100644 --- a/disco/gateway/packets.py +++ b/disco/gateway/packets.py @@ -20,21 +20,7 @@ OPCode = Enum( class Packet(TypedClass): - @classmethod - def load_json(cls, obj): - if not obj['op']: - raise Exception('Packet struct missing op key: {}'.format(obj)) - - cls = PACKETS.get(obj['op']) - - if not cls: - raise Exception('Unknown OPCode: {}'.format(obj['op'])) - - obj.update(obj['d']) - del obj['d'] - inst = cls.from_dict(obj) - inst.seq = obj['s'] - return inst + pass class DispatchPacket(Packet): diff --git a/disco/types/__init__.py b/disco/types/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/disco/types/guild.py b/disco/types/guild.py new file mode 100644 index 0000000..397ea57 --- /dev/null +++ b/disco/types/guild.py @@ -0,0 +1,7 @@ +import skema + +from disco.util.oop import TypedClass + + +class Guild(skema.Model): + id = skema.SnowflakeType() diff --git a/disco/types/user.py b/disco/types/user.py new file mode 100644 index 0000000..5411d15 --- /dev/null +++ b/disco/types/user.py @@ -0,0 +1,5 @@ +import skema + + +class User(skema.Model): + id = skema.SnowflakeType() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d17a9cd --- /dev/null +++ b/requirements.txt @@ -0,0 +1,25 @@ +backports.ssl-match-hostname==3.5.0.1 +cffi==1.8.3 +click==6.6 +cryptography==1.5 +enum34==1.1.6 +Flask==0.11.1 +gevent==1.1.2 +greenlet==0.4.10 +holster==0.0.7 +idna==2.1 +inflection==0.3.1 +ipaddress==1.0.17 +itsdangerous==0.24 +Jinja2==2.8 +MarkupSafe==0.23 +ndg-httpsclient==0.4.2 +pyasn1==0.1.9 +pycparser==2.14 +pyOpenSSL==16.1.0 +requests==2.11.1 +six==1.10.0 +skema==0.0.1 +websocket-client==0.37.0 +Werkzeug==0.11.11 +wheel==0.24.0