From 86ec5d7eca078e125837f29dcb7a92030e4ae1b4 Mon Sep 17 00:00:00 2001 From: Lorenz Diener Date: Tue, 20 Feb 2018 14:04:17 +0100 Subject: [PATCH] Re-raise Chunked Encoding Errors as Network Errors --- mastodon/streaming.py | 43 +++++++++++++++++++++++++------------------ 1 file changed, 25 insertions(+), 18 deletions(-) diff --git a/mastodon/streaming.py b/mastodon/streaming.py index d55ad54..f59b431 100644 --- a/mastodon/streaming.py +++ b/mastodon/streaming.py @@ -7,6 +7,7 @@ import json import six from mastodon import Mastodon from mastodon.Mastodon import MastodonMalformedEventError +from requests.exceptions import ChunkedEncodingError class StreamListener(object): """Callbacks for the streaming API. Create a subclass, override the on_xxx @@ -43,25 +44,31 @@ class StreamListener(object): """ event = {} line_buffer = bytearray() - for chunk in response.iter_content(chunk_size = 1): - if chunk: - if chunk == b'\n': - try: - line = line_buffer.decode('utf-8') - except UnicodeDecodeError as err: - six.raise_from( - MastodonMalformedEventError("Malformed UTF-8"), - err - ) - if line == '': - self._dispatch(event) - event = {} + try: + for chunk in response.iter_content(chunk_size = 1): + if chunk: + if chunk == b'\n': + try: + line = line_buffer.decode('utf-8') + except UnicodeDecodeError as err: + six.raise_from( + MastodonMalformedEventError("Malformed UTF-8"), + err + ) + if line == '': + self._dispatch(event) + event = {} + else: + event = self._parse_line(line, event) + line_buffer = bytearray() else: - event = self._parse_line(line, event) - line_buffer = bytearray() - else: - line_buffer.extend(chunk) - + line_buffer.extend(chunk) + except ChunkedEncodingError as err: + six.raise_from( + MastodonNetworkError("Server ceased communication."), + err + ) + def _parse_line(self, line, event): if line.startswith(':'): self.handle_heartbeat()