Add support for conversation streaming

3 files changed, 18 insertions, 1 deletions

diff --git a/docs/index.rst b/docs/index.rst
index 48fb0c1..996dee7 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -994,6 +994,7 @@ StreamListener
 .. automethod:: StreamListener.on_update
 .. automethod:: StreamListener.on_notification
 .. automethod:: StreamListener.on_delete
+.. automethod:: StreamListener.on_conversation
 .. automethod:: StreamListener.on_abort
 .. automethod:: StreamListener.handle_heartbeat
 
diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py
index edb67f8..b19c3de 100644
--- a/mastodon/Mastodon.py
+++ b/mastodon/Mastodon.py
@@ -2188,6 +2188,13 @@ class Mastodon:
         id =  self.__unpack_id(id)
         return self.__stream("/api/v1/streaming/list?list={}".format(id), listener, run_async=run_async, timeout=timeout, reconnect_async=reconnect_async, reconnect_async_wait_sec=reconnect_async_wait_sec)
     
+    @api_version("2.6.0", "2.6.0", __DICT_VERSION_STATUS)
+    def stream_direct(self, listener, run_async=False, timeout=__DEFAULT_STREAM_TIMEOUT, reconnect_async=False, reconnect_async_wait_sec=__DEFAULT_STREAM_RECONNECT_WAIT_SEC):
+        """
+        Streams direct message events for the logged-in user, as conversation events.
+        """
+        return self.__stream('/api/v1/streaming/direct', listener, run_async=run_async, timeout=timeout, reconnect_async=reconnect_async, reconnect_async_wait_sec=reconnect_async_wait_sec)
+    
     ###
     # Internal helpers, dragons probably
     ###
diff --git a/mastodon/streaming.py b/mastodon/streaming.py
index 1e1eefd..098863f 100644
--- a/mastodon/streaming.py
+++ b/mastodon/streaming.py
@@ -40,6 +40,11 @@ class StreamListener(object):
         """A status has been deleted. status_id is the status' integer ID."""
         pass
 
+    def on_conversation(self, conversation):
+        """A direct message (in the direct stream) has been received. conversation
+        contains the resulting conversation dict."""
+        pass
+
     def handle_heartbeat(self):
         """The server has sent us a keep-alive message. This callback may be
         useful to carry out periodic housekeeping tasks, or just to confirm
@@ -151,7 +156,7 @@ class CallbackStreamListener(StreamListener):
     Simple callback stream handler class.
     Can optionally additionally send local update events to a separate handler.
     """
-    def __init__(self, update_handler = None, local_update_handler = None, delete_handler = None, notification_handler = None):
+    def __init__(self, update_handler = None, local_update_handler = None, delete_handler = None, notification_handler = None, conversation_handler = None):
         super(CallbackStreamListener, self).__init__()
         self.update_handler = update_handler
         self.local_update_handler = local_update_handler
@@ -178,3 +183,7 @@ class CallbackStreamListener(StreamListener):
     def on_notification(self, notification):
         if self.notification_handler != None:
             self.notification_handler(notification)
+            
+    def on_conversation(self, conversation):
+        if self.conversation_handler != None:
+            self.conversation_handler(conversation)