diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py index a2f7738..ca4849f 100644 --- a/mastodon/Mastodon.py +++ b/mastodon/Mastodon.py @@ -1684,7 +1684,7 @@ class Mastodon: self._thread = threading.current_thread() with closing(connection) as r: try: - listener.handle_stream(r.iter_lines(chunk_size = 1, decode_unicode = True)) + listener.handle_stream(r) except AttributeError as e: if not self.closed: raise e @@ -1699,7 +1699,7 @@ class Mastodon: else: # Blocking, never returns (can only leave via exception) with closing(connection) as r: - listener.handle_stream(r.iter_lines()) + listener.handle_stream(r) def __generate_params(self, params, exclude=[]): """ diff --git a/mastodon/streaming.py b/mastodon/streaming.py index 92a02dc..3df3298 100644 --- a/mastodon/streaming.py +++ b/mastodon/streaming.py @@ -34,38 +34,46 @@ class StreamListener(object): that the connection is still open.""" pass - def handle_stream(self, lines): + def handle_stream(self, response): """ Handles a stream of events from the Mastodon server. When each event is received, the corresponding .on_[name]() method is called. - lines: an iterable of lines of bytes sent by the Mastodon server, as - returned by requests.Response.iter_lines(). + response; a requests response object with the open stream for reading. """ - event = {} - for raw_line in lines: - try: - line = raw_line.decode('utf-8') - except UnicodeDecodeError as err: - six.raise_from( - MastodonMalformedEventError("Malformed UTF-8", line), - err - ) - - if line.startswith(':'): - self.handle_heartbeat() - elif line == '': - # end of event - self._dispatch(event) - event = {} - else: - key, value = line.split(': ', 1) - # According to the MDN spec, repeating the 'data' key - # represents a newline(!) - if key in event: - event[key] += '\n' + value + self.event = {} + line_buffer = bytearray() + for chunk in response.iter_content(chunk_size = 1): + if chunk: + if chunk == b'\n': + self.handle_line(line_buffer) + line_buffer = bytearray() else: - event[key] = value + line_buffer.extend(chunk) + + def handle_line(self, raw_line): + try: + line = raw_line.decode('utf-8') + except UnicodeDecodeError as err: + six.raise_from( + MastodonMalformedEventError("Malformed UTF-8", line), + err + ) + + if line.startswith(':'): + self.handle_heartbeat() + elif line == '': + # end of event + self._dispatch(self.event) + self.event = {} + else: + key, value = line.split(': ', 1) + # According to the MDN spec, repeating the 'data' key + # represents a newline(!) + if key in self.event: + self.event[key] += '\n' + value + else: + self.event[key] = value def _dispatch(self, event): try: