From 0e3f518ad1f58329ed20fc983a5fa8d5691bc90a Mon Sep 17 00:00:00 2001
From: Andrei <b1naryth1ef@gmail.com>
Date: Wed, 19 Oct 2016 20:15:00 -0500
Subject: [PATCH] Add Client.update_presence, fix heartbeating on dead
 connections

---
 disco/client.py         | 22 ++++++++++++++++++++++
 disco/gateway/client.py |  6 ++++++
 disco/types/base.py     | 15 +++++++++++++--
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/disco/client.py b/disco/client.py
index 0c15c50..d447ecc 100644
--- a/disco/client.py
+++ b/disco/client.py
@@ -1,3 +1,4 @@
+import time
 import gevent
 
 from holster.emitter import Emitter
@@ -5,6 +6,8 @@ from holster.emitter import Emitter
 from disco.state import State, StateConfig
 from disco.api.client import APIClient
 from disco.gateway.client import GatewayClient
+from disco.gateway.packets import OPCode
+from disco.types.user import Status, Game
 from disco.util.config import Config
 from disco.util.logging import LoggingClass
 from disco.util.backdoor import DiscoBackdoorServer
@@ -99,6 +102,25 @@ class Client(object):
                 localf=lambda: self.manhole_locals)
             self.manhole.start()
 
+    def update_presence(self, game=None, status=None, afk=False, since=0.0):
+        if game and not isinstance(game, Game):
+            raise TypeError('Game must be a Game model')
+
+        if status is Status.IDLE and not since:
+            since = int(time.time() * 1000)
+
+        payload = {
+            'afk': afk,
+            'since': since,
+            'status': status.value,
+            'game': None,
+        }
+
+        if game:
+            payload['game'] = game.to_dict()
+
+        self.gw.send(OPCode.STATUS_UPDATE, payload)
+
     def run(self):
         """
         Run the client (e.g. the :class:`GatewayClient`) in a new greenlet
diff --git a/disco/gateway/client.py b/disco/gateway/client.py
index 61f7dfd..251b800 100644
--- a/disco/gateway/client.py
+++ b/disco/gateway/client.py
@@ -173,10 +173,16 @@ class GatewayClient(LoggingClass):
             })
 
     def on_close(self, code, reason):
+        # Kill heartbeater, a reconnect/resume will trigger a HELLO which will
+        #  respawn it
+        self._heartbeat_task.kill()
+
+        # If we're quitting, just break out of here
         if self.shutting_down:
             self.log.info('WS Closed: shutting down')
             return
 
+        # Track reconnect attempts
         self.reconnects += 1
         self.log.info('WS Closed: [%s] %s (%s)', code, reason, self.reconnects)
 
diff --git a/disco/types/base.py b/disco/types/base.py
index 63bc628..78536a3 100644
--- a/disco/types/base.py
+++ b/disco/types/base.py
@@ -3,7 +3,7 @@ import gevent
 import inspect
 import functools
 
-from holster.enum import BaseEnumMeta
+from holster.enum import BaseEnumMeta, EnumAttr
 from datetime import datetime as real_datetime
 
 from disco.util.functional import CachedSlotProperty
@@ -33,6 +33,14 @@ class FieldType(object):
         else:
             self.typ = lambda raw, _: typ(raw)
 
+    def serialize(self, value):
+        if isinstance(value, EnumAttr):
+            return value.value
+        elif isinstance(value, Model):
+            return value.to_dict()
+        else:
+            return value
+
     def try_convert(self, raw, client):
         pass
 
@@ -260,7 +268,10 @@ class Model(six.with_metaclass(ModelMeta, AsyncChainable)):
                     pass
 
     def to_dict(self):
-        return {k: getattr(self, k) for k in six.iterkeys(self.__class__._fields)}
+        obj = {}
+        for name, field in six.iteritems(self.__class__._fields):
+            obj[name] = field.serialize(getattr(self, name))
+        return obj
 
     @classmethod
     def create(cls, client, data, **kwargs):