Fix streaming API to be more stable (closes #117)

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=[]):
         """

mastodon/streaming.py

@@ -34,16 +34,24 @@ 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:
+        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:
+                    line_buffer.extend(chunk)
+        
+    def handle_line(self, raw_line):
         try:
             line = raw_line.decode('utf-8')
         except UnicodeDecodeError as err:
@@ -56,16 +64,16 @@ class StreamListener(object):
             self.handle_heartbeat()
         elif line == '':
             # end of event
-                self._dispatch(event)
-                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 event:
-                    event[key] += '\n' + value
+            if key in self.event:
+                self.event[key] += '\n' + value
             else:
-                    event[key] = value
+                self.event[key] = value
 
     def _dispatch(self, event):
         try: