diff --git a/sseclient/__init__.py b/sseclient/__init__.py index 56171f9..2da259f 100644 --- a/sseclient/__init__.py +++ b/sseclient/__init__.py @@ -20,7 +20,7 @@ class SSEClient(object): specification. """ - def __init__(self, event_source, char_enc='utf-8'): + def __init__(self, event_source, char_enc='utf-8', ignore_comments=True): """Initialize the SSE client over an existing, ready to consume event source. @@ -33,6 +33,7 @@ def __init__(self, event_source, char_enc='utf-8'): event_source) self._event_source = event_source self._char_enc = char_enc + self._ignore_comments = ignore_comments def _read(self): """Read the incoming event source stream and yield event chunks. @@ -43,6 +44,12 @@ def _read(self): SSE delimiter (empty new line) to yield full, correct event chunks.""" data = b'' for chunk in self._event_source: + + # yield comments + if not self._ignore_comments and chunk.startswith(b':') and chunk.endswith(b'\n'): + yield chunk + continue + for line in chunk.splitlines(True): data += line if data.endswith((b'\r\r', b'\n\n', b'\r\n\r\n')): @@ -54,6 +61,16 @@ def _read(self): def events(self): for chunk in self._read(): event = Event() + + if not self._ignore_comments and chunk.startswith(b':') and chunk.endswith(b'\n'): + event.event = 'comment' + event.data = chunk[1:-2].decode(self._char_enc) + + # Dispatch the event + self._logger.debug('Dispatching %s...', event) + yield event + continue + # Split before decoding so splitlines() only uses \r and \n for line in chunk.splitlines(): # Decode the line. diff --git a/tests/unittests.py b/tests/unittests.py index d19aa32..432aad2 100644 --- a/tests/unittests.py +++ b/tests/unittests.py @@ -15,9 +15,9 @@ import sseclient # noqa: E402 -def parse(content): +def parse(content, ignore_comments=True): return [{'id': ev.id, 'event': ev.event, 'data': ev.data} - for ev in sseclient.SSEClient(content).events()] + for ev in sseclient.SSEClient(content, ignore_comments=ignore_comments).events()] class Parser(unittest.TestCase): @@ -104,12 +104,27 @@ def test_ignores_comments(self): [{'id': None, 'event': 'message', 'data': 'Hello'}, {'id': None, 'event': 'message', 'data': 'World'}]) + def test_do_not_ignore_comments(self): + self.assertEqual( + parse([(b'data: Hello\n\n:nothing to see here\n\n' + b'data: World\n\n')], ignore_comments=False), + [{'id': None, 'event': 'message', 'data': 'Hello'}, + {'id': None, 'event': 'comment', 'data': 'nothing to see here'}, + {'id': None, 'event': 'message', 'data': 'World'}]) + def test_ignores_empty_comments(self): self.assertEqual( parse([b'data: Hello\n\n:\n\ndata: World\n\n']), [{'id': None, 'event': 'message', 'data': 'Hello'}, {'id': None, 'event': 'message', 'data': 'World'}]) + def test_do_not_ignore_empty_comments(self): + self.assertEqual( + parse([b'data: Hello\n\n:\n\ndata: World\n\n'], ignore_comments=False), + [{'id': None, 'event': 'message', 'data': 'Hello'}, + {'id': None, 'event': 'comment', 'data': ''}, + {'id': None, 'event': 'message', 'data': 'World'}]) + def test_does_not_ignore_multiline_strings(self): self.assertEqual( parse([b'data: line one\ndata:\ndata: line two\n\n']),