diff --git a/stream/ws_client.py b/stream/ws_client.py index 1cc56cdd..d3da3765 100644 --- a/stream/ws_client.py +++ b/stream/ws_client.py @@ -41,7 +41,10 @@ def __init__(self, configuration, url, headers): header = [] self._connected = False self._channels = {} - self._all = "" + if six.PY3: + self._all = b"" + else: + self._all = "" # We just need to pass the Authorization, ignore all the other # http headers we get from the generated code @@ -77,7 +80,11 @@ def peek_channel(self, channel, timeout=0): empty string otherwise.""" self.update(timeout=timeout) if channel in self._channels: + if six.PY3: + return self._channels[channel].decode("utf-8", "replace") return self._channels[channel] + if six.PY3: + return b"" return "" def read_channel(self, channel, timeout=0): @@ -88,6 +95,8 @@ def read_channel(self, channel, timeout=0): ret = self._channels[channel] if channel in self._channels: del self._channels[channel] + if six.PY3 and isinstance(ret, bytes): + return ret.decode("utf-8", "replace") return ret def readline_channel(self, channel, timeout=None): @@ -98,15 +107,21 @@ def readline_channel(self, channel, timeout=None): while self.is_open() and time.time() - start < timeout: if channel in self._channels: data = self._channels[channel] - if "\n" in data: - index = data.find("\n") + newline_symbol = "\n" + if six.PY3: + newline_Symbol = b"\n" + if newline_symbol in data: + index = data.find(newline_symbol) ret = data[:index] data = data[index+1:] if data: self._channels[channel] = data else: del self._channels[channel] - return ret + if six.PY3 and isinstance(ret, bytes): + return ret.decode("utf-8", "replace") + else: + return ret self.update(timeout=(timeout - time.time() + start)) def write_channel(self, channel, data): @@ -147,9 +162,15 @@ def read_all(self): channels mapped for each input. """ out = self._all - self._all = "" + if six.PY3: + self._all = b"" + else: + self._all = "" self._channels = {} - return out + if six.PY3: + return out.decode("utf-8", "replace") + else: + return out def is_open(self): """True if the connection is still alive.""" @@ -175,10 +196,11 @@ def update(self, timeout=0): return elif op_code == ABNF.OPCODE_BINARY or op_code == ABNF.OPCODE_TEXT: data = frame.data - if six.PY3: - data = data.decode("utf-8") if len(data) > 1: - channel = ord(data[0]) + if six.PY3: + channel = data[0] + else: + channel = ord(data[0]) data = data[1:] if data: if channel in [STDOUT_CHANNEL, STDERR_CHANNEL]: