Many fixes for streaming stuff

3 files changed, 42 insertions, 32 deletions

diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py
index c35ad2d..9b28861 100644
--- a/mastodon/Mastodon.py
+++ b/mastodon/Mastodon.py
@@ -17,10 +17,11 @@ import re
 import copy
 import threading
 import sys
+
 try:
     from urllib.parse import urlparse
 except ImportError:
-     from urlparse import urlparse
+    from urlparse import urlparse
 
 
 class Mastodon:
@@ -1014,12 +1015,6 @@ class Mastodon:
         Streams events that are relevant to the authorized user, i.e. home
         timeline and notifications. 'listener' should be a subclass of
         StreamListener which will receive callbacks for incoming events.
-
-        If async is False, this method blocks forever.
-
-        If async is True, 'listener' will listen on another thread and this method
-        will return a handle corresponding to the open connection. The
-        connection may be closed at any time by calling its close() method.
         """
         return self.__stream('/api/v1/streaming/user', listener, async=async)
 
@@ -1027,12 +1022,6 @@ class Mastodon:
         """
         Streams public events. 'listener' should be a subclass of StreamListener
         which will receive callbacks for incoming events.
-
-        If async is False, this method blocks forever.
-
-        If async is True, 'listener' will listen on another thread and this method
-        will return a handle corresponding to the open connection. The
-        connection may be closed at any time by calling its close() method.
         """
         return self.__stream('/api/v1/streaming/public', listener, async=async)
 
@@ -1041,11 +1030,6 @@ class Mastodon:
         Streams local events. 'listener' should be a subclass of StreamListener
         which will receive callbacks for incoming events.
 
-        If async is False, this method blocks forever.
-
-        If async is True, 'listener' will listen on another thread and this method
-        will return a handle corresponding to the open connection. The
-        connection may be closed at any time by calling its close() method.
         """
         return self.__stream('/api/v1/streaming/public/local', listener, async=async)
 
@@ -1054,12 +1038,6 @@ class Mastodon:
         Returns all public statuses for the hashtag 'tag'. 'listener' should be
         a subclass of StreamListener which will receive callbacks for incoming
         events.
-
-        If async is False, this method blocks forever.
-
-        If async is True, 'listener' will listen on another thread and this method
-        will return a handle corresponding to the open connection. The
-        connection may be closed at any time by calling its close() method.
         """
         return self.__stream("/api/v1/streaming/hashtag?tag={}".format(tag), listener)
 
diff --git a/mastodon/__init__.py b/mastodon/__init__.py
index 3123356..fdf776d 100644
--- a/mastodon/__init__.py
+++ b/mastodon/__init__.py
@@ -1,4 +1,4 @@
 from mastodon.Mastodon import Mastodon
-from mastodon.streaming import StreamListener
+from mastodon.streaming import StreamListener, CallbackStreamListener
 
-__all__ = ['Mastodon', 'StreamListener']
+__all__ = ['Mastodon', 'StreamListener', 'CallbackStreamListener']
diff --git a/mastodon/streaming.py b/mastodon/streaming.py
index cac2456..92a02dc 100644
--- a/mastodon/streaming.py
+++ b/mastodon/streaming.py
@@ -67,10 +67,6 @@ class StreamListener(object):
                 else:
                     event[key] = value
 
-        # end of stream
-        if event:
-            log.warn("outstanding partial event at end of stream: %s", event)
-
     def _dispatch(self, event):
         try:
             name = event['event']
@@ -92,8 +88,44 @@ class StreamListener(object):
         handler_name = 'on_' + name
         try:
             handler = getattr(self, handler_name)
-        except AttributeError:
-            log.warn("Unhandled event '%s'", name)
+        except AttributeError as err:
+            six.raise_from(
+               MastodonMalformedEventError('Bad event type', name),
+               err
+            )
         else:
             # TODO: allow handlers to return/raise to stop streaming cleanly
             handler(payload)
+
+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):
+        super(CallbackStreamListener, self).__init__()
+        self.update_handler = update_handler
+        self.local_update_handler = local_update_handler
+        self.delete_handler = delete_handler
+        self.notification_handler = notification_handler
+        
+    def on_update(self, status):
+        if self.update_handler != None:
+            self.update_handler(status)
+        
+        try:
+            if self.local_update_handler != None and not "@" in status["account"]["acct"]:
+                self.local_update_handler(status)
+        except Exception as err:
+            six.raise_from(
+               MastodonMalformedEventError('received bad update', status),
+               err
+            )
+    
+    def on_delete(self, deleted_id):
+        if self.delete_handler != None:
+            self.delete_handler(deleted_id)
+    
+    def on_notification(self, notification):
+        if self.notification_handler != None:
+            self.notification_handler(notification)
\ No newline at end of file