import re import time import warnings import threading import six import requests # Technically, we should support streams that mix line endings. This regex, # however, assumes that a system will provide consistent line endings. end_of_field = re.compile(r'\r\n\r\n|\r\r|\n\n') class SSEClient(object): def __init__(self, url, session, build_headers, last_id=None, retry=3000, **kwargs): self.url = url self.last_id = last_id self.retry = retry self.running = True # Optional support for passing in a requests.Session() self.session = session # function for building auth header when token expires self.build_headers = build_headers self.start_time = None # Any extra kwargs will be fed into the requests.get call later. self.requests_kwargs = kwargs # The SSE spec requires making requests with Cache-Control: nocache if 'headers' not in self.requests_kwargs: self.requests_kwargs['headers'] = {} self.requests_kwargs['headers']['Cache-Control'] = 'no-cache' # The 'Accept' header is not required, but explicit > implicit self.requests_kwargs['headers']['Accept'] = 'text/event-stream' # Keep data here as it streams in self.buf = u'' self._connect() def _connect(self): if self.last_id: self.requests_kwargs['headers']['Last-Event-ID'] = self.last_id headers = self.build_headers() self.requests_kwargs['headers'].update(headers) # Use session if set. Otherwise fall back to requests module. self.requester = self.session or requests self.resp = self.requester.get(self.url, stream=True, **self.requests_kwargs) self.resp_iterator = self.resp.iter_content(decode_unicode=True) # TODO: Ensure we're handling redirects. Might also stick the 'origin' # attribute on Events like the Javascript spec requires. self.resp.raise_for_status() def _event_complete(self): return re.search(end_of_field, self.buf) is not None def __iter__(self): return self def __next__(self): while not self._event_complete(): try: nextchar = next(self.resp_iterator) self.buf += nextchar except (StopIteration, requests.RequestException): time.sleep(self.retry / 1000.0) self._connect() # The SSE spec only supports resuming from a whole message, so # if we have half a message we should throw it out. head, sep, tail = self.buf.rpartition('\n') self.buf = head + sep continue split = re.split(end_of_field, self.buf) head = split[0] tail = "".join(split[1:]) self.buf = tail msg = Event.parse(head) if msg.data == "credential is no longer valid": self._connect() return None if msg.data == 'null': return None # If the server requests a specific retry delay, we need to honor it. if msg.retry: self.retry = msg.retry # last_id should only be set if included in the message. It's not # forgotten if a message omits it. if msg.id: self.last_id = msg.id return msg if six.PY2: next = __next__ class Event(object): sse_line_pattern = re.compile('(?P[^:]*):?( ?(?P.*))?') def __init__(self, data='', event='message', id=None, retry=None): self.data = data self.event = event self.id = id self.retry = retry def dump(self): lines = [] if self.id: lines.append('id: %s' % self.id) # Only include an event line if it's not the default already. if self.event != 'message': lines.append('event: %s' % self.event) if self.retry: lines.append('retry: %s' % self.retry) lines.extend('data: %s' % d for d in self.data.split('\n')) return '\n'.join(lines) + '\n\n' @classmethod def parse(cls, raw): """ Given a possibly-multiline string representing an SSE message, parse it and return a Event object. """ msg = cls() for line in raw.split('\n'): m = cls.sse_line_pattern.match(line) if m is None: # Malformed line. Discard but warn. warnings.warn('Invalid SSE line: "%s"' % line, SyntaxWarning) continue name = m.groupdict()['name'] value = m.groupdict()['value'] if name == '': # line began with a ":", so is a comment. Ignore continue if name == 'data': # If we already have some data, then join to it with a newline. # Else this is it. if msg.data: msg.data = '%s\n%s' % (msg.data, value) else: msg.data = value elif name == 'event': msg.event = value elif name == 'id': msg.id = value elif name == 'retry': msg.retry = int(value) return msg def __str__(self): return self.data