diff options
Diffstat (limited to 'mastodon/streaming.py')
-rw-r--r-- | mastodon/streaming.py | 91 |
1 files changed, 58 insertions, 33 deletions
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 | |||
4 | """ | 4 | """ |
5 | 5 | ||
6 | import json | 6 | import json |
7 | import logging | ||
8 | import six | 7 | import six |
9 | 8 | from mastodon import Mastodon | |
10 | 9 | from mastodon.Mastodon import MastodonMalformedEventError | |
11 | log = logging.getLogger(__name__) | ||
12 | |||
13 | |||
14 | class MalformedEventError(Exception): | ||
15 | """Raised when the server-sent event stream is malformed.""" | ||
16 | pass | ||
17 | |||
18 | 10 | ||
19 | class StreamListener(object): | 11 | class StreamListener(object): |
20 | """Callbacks for the streaming API. Create a subclass, override the on_xxx | 12 | """Callbacks for the streaming API. Create a subclass, override the on_xxx |
@@ -24,7 +16,7 @@ class StreamListener(object): | |||
24 | 16 | ||
25 | def on_update(self, status): | 17 | def on_update(self, status): |
26 | """A new status has appeared! 'status' is the parsed JSON dictionary | 18 | """A new status has appeared! 'status' is the parsed JSON dictionary |
27 | describing the status.""" | 19 | describing the status.""" |
28 | pass | 20 | pass |
29 | 21 | ||
30 | def on_notification(self, notification): | 22 | def on_notification(self, notification): |
@@ -40,7 +32,8 @@ describing the status.""" | |||
40 | """The server has sent us a keep-alive message. This callback may be | 32 | """The server has sent us a keep-alive message. This callback may be |
41 | useful to carry out periodic housekeeping tasks, or just to confirm | 33 | useful to carry out periodic housekeeping tasks, or just to confirm |
42 | that the connection is still open.""" | 34 | that the connection is still open.""" |
43 | 35 | pass | |
36 | |||
44 | def handle_stream(self, lines): | 37 | def handle_stream(self, lines): |
45 | """ | 38 | """ |
46 | Handles a stream of events from the Mastodon server. When each event | 39 | Handles a stream of events from the Mastodon server. When each event |
@@ -55,7 +48,7 @@ describing the status.""" | |||
55 | line = raw_line.decode('utf-8') | 48 | line = raw_line.decode('utf-8') |
56 | except UnicodeDecodeError as err: | 49 | except UnicodeDecodeError as err: |
57 | six.raise_from( | 50 | six.raise_from( |
58 | MalformedEventError("Malformed UTF-8", line), | 51 | MastodonMalformedEventError("Malformed UTF-8", line), |
59 | err | 52 | err |
60 | ) | 53 | ) |
61 | 54 | ||
@@ -63,7 +56,7 @@ describing the status.""" | |||
63 | self.handle_heartbeat() | 56 | self.handle_heartbeat() |
64 | elif line == '': | 57 | elif line == '': |
65 | # end of event | 58 | # end of event |
66 | self._despatch(event) | 59 | self._dispatch(event) |
67 | event = {} | 60 | event = {} |
68 | else: | 61 | else: |
69 | key, value = line.split(': ', 1) | 62 | key, value = line.split(': ', 1) |
@@ -74,33 +67,65 @@ describing the status.""" | |||
74 | else: | 67 | else: |
75 | event[key] = value | 68 | event[key] = value |
76 | 69 | ||
77 | # end of stream | 70 | def _dispatch(self, event): |
78 | if event: | ||
79 | log.warn("outstanding partial event at end of stream: %s", event) | ||
80 | |||
81 | def _despatch(self, event): | ||
82 | try: | 71 | try: |
83 | name = event['event'] | 72 | name = event['event'] |
84 | data = event['data'] | 73 | data = event['data'] |
85 | payload = json.loads(data) | 74 | payload = json.loads(data, object_hook = Mastodon._Mastodon__json_hooks) |
86 | except KeyError as err: | 75 | except KeyError as err: |
87 | six.raise_from( | 76 | six.raise_from( |
88 | MalformedEventError('Missing field', err.args[0], event), | 77 | MastodonMalformedEventError('Missing field', err.args[0], event), |
89 | err | 78 | err |
90 | ) | 79 | ) |
91 | except ValueError as err: | 80 | except ValueError as err: |
92 | # py2: plain ValueError | 81 | # py2: plain ValueError |
93 | # py3: json.JSONDecodeError, a subclass of ValueError | 82 | # py3: json.JSONDecodeError, a subclass of ValueError |
94 | six.raise_from( | 83 | six.raise_from( |
95 | MalformedEventError('Bad JSON', data), | 84 | MastodonMalformedEventError('Bad JSON', data), |
96 | err | 85 | err |
97 | ) | 86 | ) |
98 | 87 | ||
99 | handler_name = 'on_' + name | 88 | handler_name = 'on_' + name |
100 | try: | 89 | try: |
101 | handler = getattr(self, handler_name) | 90 | handler = getattr(self, handler_name) |
102 | except AttributeError: | 91 | except AttributeError as err: |
103 | log.warn("Unhandled event '%s'", name) | 92 | six.raise_from( |
93 | MastodonMalformedEventError('Bad event type', name), | ||
94 | err | ||
95 | ) | ||
104 | else: | 96 | else: |
105 | # TODO: allow handlers to return/raise to stop streaming cleanly | 97 | # TODO: allow handlers to return/raise to stop streaming cleanly |
106 | handler(payload) | 98 | handler(payload) |
99 | |||
100 | class CallbackStreamListener(StreamListener): | ||
101 | """ | ||
102 | Simple callback stream handler class. | ||
103 | Can optionally additionally send local update events to a separate handler. | ||
104 | """ | ||
105 | def __init__(self, update_handler = None, local_update_handler = None, delete_handler = None, notification_handler = None): | ||
106 | super(CallbackStreamListener, self).__init__() | ||
107 | self.update_handler = update_handler | ||
108 | self.local_update_handler = local_update_handler | ||
109 | self.delete_handler = delete_handler | ||
110 | self.notification_handler = notification_handler | ||
111 | |||
112 | def on_update(self, status): | ||
113 | if self.update_handler != None: | ||
114 | self.update_handler(status) | ||
115 | |||
116 | try: | ||
117 | if self.local_update_handler != None and not "@" in status["account"]["acct"]: | ||
118 | self.local_update_handler(status) | ||
119 | except Exception as err: | ||
120 | six.raise_from( | ||
121 | MastodonMalformedEventError('received bad update', status), | ||
122 | err | ||
123 | ) | ||
124 | |||
125 | def on_delete(self, deleted_id): | ||
126 | if self.delete_handler != None: | ||
127 | self.delete_handler(deleted_id) | ||
128 | |||
129 | def on_notification(self, notification): | ||
130 | if self.notification_handler != None: | ||
131 | self.notification_handler(notification) \ No newline at end of file | ||