diff --git a/i3.py b/i3.py index 48555d2..1020e6a 100644 --- a/i3.py +++ b/i3.py @@ -127,30 +127,25 @@ class Socket(object): - path of the i3 socket. Path is retrieved from i3-wm itself via "i3.get_socket_path()" if not provided. - timeout in seconds - - chunk_size in bytes - magic_string as a safety string for i3-ipc. Set to 'i3-ipc' by default. """ - magic_string = 'i3-ipc' # safety string for i3-ipc - chunk_size = 1024 # in bytes + magic_string = str.encode('i3-ipc') # safety string for i3-ipc timeout = 0.5 # in seconds buffer = b'' # byte string - def __init__(self, path=None, timeout=None, chunk_size=None, - magic_string=None): + def __init__(self, path=None, timeout=None, magic_string=None): if not path: path = get_socket_path() self.path = path if timeout: self.timeout = timeout - if chunk_size: - self.chunk_size = chunk_size if magic_string: self.magic_string = magic_string # Socket initialization and connection self.initialize() self.connect() # Struct format initialization, length of magic string is in bytes - self.struct_header = '<%dsII' % len(self.magic_string.encode('utf-8')) + self.struct_header = '<%dsII' % len(self.magic_string) self.struct_header_size = struct.calcsize(self.struct_header) def initialize(self): @@ -173,7 +168,7 @@ def connect(self, path=None): except socket.error: raise ConnectionError(path) - def get(self, msg_type, payload=''): + def get(self, msg_type, payload=str.encode('')): """ Convenience method, calls "socket.send(msg_type, payload)" and returns data from "socket.receive()". @@ -190,10 +185,10 @@ def subscribe(self, event_type, event=None): payload = [event_type] if event: payload.append(event) - payload = json.dumps(payload) + payload = str.encode(json.dumps(payload)) return self.get('subscribe', payload) - def send(self, msg_type, payload=''): + def send(self, msg_type, payload=str.encode('')): """ Sends the given message type with given message by packing them and continuously sending bytes from the packed message. @@ -208,12 +203,17 @@ def receive(self): successful. Returns None on failure. """ try: - data = self.socket.recv(self.chunk_size) + data = self.socket.recv(self.struct_header_size) msg_magic, msg_length, msg_type = self.unpack_header(data) + + # Sanity check + assert msg_magic == self.magic_string \ + , "Incorrect magic string received!" + msg_size = self.struct_header_size + msg_length # Keep receiving data until the whole message gets through while len(data) < msg_size: - data += self.socket.recv(msg_length) + data += self.socket.recv(msg_size - len(data)) data = self.buffer + data return self.unpack(data) except socket.timeout: @@ -226,14 +226,14 @@ def pack(self, msg_type, payload): """ msg_magic = self.magic_string # Get the byte count instead of number of characters - msg_length = len(payload.encode('utf-8')) + msg_length = len(payload) msg_type = parse_msg_type(msg_type) - # "struct.pack" returns byte string, decoding it for concatenation - msg_length = struct.pack('I', msg_length).decode('utf-8') - msg_type = struct.pack('I', msg_type).decode('utf-8') - message = '%s%s%s%s' % (msg_magic, msg_length, msg_type, payload) + # "struct.pack" returns byte string + msg_length = struct.pack('I', msg_length) + msg_type = struct.pack('I', msg_type) + message = msg_magic + msg_length + msg_type + payload # Encoding the message back to byte string - return message.encode('utf-8') + return message def unpack(self, data): """ @@ -245,8 +245,8 @@ def unpack(self, data): msg_size = self.struct_header_size + msg_length # Message shouldn't be any longer than the data if data_size >= msg_size: - payload = data[self.struct_header_size:msg_size].decode('utf-8') - payload = json.loads(payload) + payload = data[self.struct_header_size:msg_size] + payload = json.loads(payload.decode('utf-8')) self.buffer = data[msg_size:] return payload else: @@ -365,7 +365,7 @@ def __call_cmd__(cmd): output = subprocess.check_output(cmd) except subprocess.CalledProcessError as error: output = error.output - output = output.decode('utf-8') # byte string decoding + output = output # byte string decoding return output.strip() @@ -384,7 +384,7 @@ def default_socket(socket=None): return __socket__ -def msg(type, message=''): +def msg(type, message=str.encode('')): """ Takes a message type and a message itself. Talks to the i3 via socket and returns the response from the socket. @@ -406,7 +406,7 @@ def function(*args2, **crit2): criteria.update(crit2) if criteria: msg_full = '%s %s' % (container(**criteria), msg_full) - response = msg(type, msg_full) + response = msg(type, str.encode(msg_full)) response = success(response) if isinstance(response, i3Exception): raise response