3 files changed, 84 insertions, 70 deletions

diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py
index d2d634e..88f4906 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:
@@ -1009,57 +1010,34 @@ class Mastodon:
     ###
     # Streaming
     ###
-    def user_stream(self, listener, async=False):
+    def stream_user(self, listener, async=False):
         """
         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)
 
-    def public_stream(self, listener, async=False):
+    def stream_public(self, listener, async=False):
         """
         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)
 
-    def local_stream(self, listener, async=False):
+    def stream_local(self, listener, async=False):
         """
         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)
 
-    def hashtag_stream(self, tag, listener, async=False):
+    def stream_hashtag(self, tag, listener, async=False):
         """
         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)
 
@@ -1083,8 +1061,8 @@ class Mastodon:
 
         return (date_time_utc - epoch_utc).total_seconds()
 
-
-    def __json_date_parse(self, json_object):
+    @staticmethod
+    def __json_date_parse(json_object):
         """
         Parse dates in certain known json fields, if possible.
         """
@@ -1100,7 +1078,8 @@ class Mastodon:
                     raise MastodonAPIError('Encountered invalid date.')
         return json_object
 
-    def __json_id_to_bignum(self, json_object):
+    @staticmethod
+    def __json_id_to_bignum(json_object):
         """
         Converts json string IDs to native python bignums.
         """
@@ -1117,10 +1096,11 @@ class Mastodon:
                 pass
 
         return json_object
-
-    def __json_hooks(self, json_object):
-        json_object = self.__json_date_parse(json_object)
-        json_object = self.__json_id_to_bignum(json_object)
+    
+    @staticmethod
+    def __json_hooks(json_object):
+        json_object = Mastodon.__json_date_parse(json_object)
+        json_object = Mastodon.__json_id_to_bignum(json_object)
         return json_object
 
     def __api_request(self, method, endpoint, params={}, files={}, do_ratelimiting=True):
@@ -1424,6 +1404,7 @@ class MastodonError(Exception):
 
 
 class MastodonIllegalArgumentError(ValueError, MastodonError):
+    """Raised when an incorrect parameter is passed to a function"""
     pass
 
 
@@ -1432,16 +1413,24 @@ class MastodonIOError(IOError, MastodonError):
 
 
 class MastodonFileNotFoundError(MastodonIOError):
+    """Raised when a file requested to be loaded can not be opened"""
     pass
 
 
 class MastodonNetworkError(MastodonIOError):
+    """Raised when network communication with the server fails"""
     pass
 
 
 class MastodonAPIError(MastodonError):
+    """Raised when the mastodon API generates a response that cannot be handled"""
     pass
 
 
 class MastodonRatelimitError(MastodonError):
+    """Raised when rate limiting is set to manual mode and the rate limit is exceeded"""
+    pass
+
+class MastodonMalformedEventError(MastodonError):
+    """Raised when the server-sent event stream is malformed"""
     pass
diff --git a/mastodon/__init__.py b/mastodon/__init__.py
index 9c8e39b..fdf776d 100644
--- a/mastodon/__init__.py
+++ b/mastodon/__init__.py
@@ -1,4 +1,4 @@
 from mastodon.Mastodon import Mastodon
-from mastodon.streaming import StreamListener, MalformedEventError
+from mastodon.streaming import StreamListener, CallbackStreamListener
 
-__all__ = ['Mastodon', 'StreamListener', 'MalformedEventError']
+__all__ = ['Mastodon', 'StreamListener', 'CallbackStreamListener']
diff --git a/mastodon/streaming.py b/mastodon/streaming.py
index 290ed44..92a02dc 100644
--- a/mastodon/streaming.py
+++ b/mastodon/streaming.py
@@ -4,17 +4,9 @@ https://github.com/tootsuite/mastodon/blob/master/docs/Using-the-API/Streaming-A
 """
 
 import json
-import logging
 import six
-
-
-log = logging.getLogger(__name__)
-
-
-class MalformedEventError(Exception):
-    """Raised when the server-sent event stream is malformed."""
-    pass
-
+from mastodon import Mastodon
+from mastodon.Mastodon import MastodonMalformedEventError
 
 class StreamListener(object):
     """Callbacks for the streaming API. Create a subclass, override the on_xxx
@@ -24,7 +16,7 @@ class StreamListener(object):
 
     def on_update(self, status):
         """A new status has appeared! 'status' is the parsed JSON dictionary
-describing the status."""
+        describing the status."""
         pass
 
     def on_notification(self, notification):
@@ -40,7 +32,8 @@ describing the status."""
         """The server has sent us a keep-alive message. This callback may be
         useful to carry out periodic housekeeping tasks, or just to confirm
         that the connection is still open."""
-
+        pass
+    
     def handle_stream(self, lines):
         """
         Handles a stream of events from the Mastodon server. When each event
@@ -55,7 +48,7 @@ describing the status."""
                 line = raw_line.decode('utf-8')
             except UnicodeDecodeError as err:
                 six.raise_from(
-                    MalformedEventError("Malformed UTF-8", line),
+                    MastodonMalformedEventError("Malformed UTF-8", line),
                     err
                 )
 
@@ -63,7 +56,7 @@ describing the status."""
                 self.handle_heartbeat()
             elif line == '':
                 # end of event
-                self._despatch(event)
+                self._dispatch(event)
                 event = {}
             else:
                 key, value = line.split(': ', 1)
@@ -74,33 +67,65 @@ describing the status."""
                 else:
                     event[key] = value
 
-        # end of stream
-        if event:
-            log.warn("outstanding partial event at end of stream: %s", event)
-
-    def _despatch(self, event):
+    def _dispatch(self, event):
         try:
             name = event['event']
             data = event['data']
-            payload = json.loads(data)
+            payload = json.loads(data, object_hook = Mastodon._Mastodon__json_hooks)
         except KeyError as err:
-            six.raise_from(
-                MalformedEventError('Missing field', err.args[0], event),
-                err
-            )
+           six.raise_from(
+               MastodonMalformedEventError('Missing field', err.args[0], event),
+               err
+           )
         except ValueError as err:
-            # py2: plain ValueError
-            # py3: json.JSONDecodeError, a subclass of ValueError
-            six.raise_from(
-                MalformedEventError('Bad JSON', data),
-                err
-            )
-
+           # py2: plain ValueError
+           # py3: json.JSONDecodeError, a subclass of ValueError
+           six.raise_from(
+               MastodonMalformedEventError('Bad JSON', data),
+               err
+           )
+           
         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