path: root/.i3/scripts/i3.py

#======================================================================
# i3 (Python module for communicating with i3 window manager)
# Copyright (C) 2012  Jure Ziberna
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#======================================================================


import sys
import subprocess
import json
import socket
import struct
import threading
import time

ModuleType = type(sys)


__author__ = 'Jure Ziberna'
__version__ = '0.6.5'
__date__ = '2012-06-20'
__license__ = 'GNU GPL 3'


MSG_TYPES = [
    'command',
    'get_workspaces',
    'subscribe',
    'get_outputs',
    'get_tree',
    'get_marks',
    'get_bar_config',
]

EVENT_TYPES = [
    'workspace',
    'output',
]


class i3Exception(Exception):
    pass

class MessageTypeError(i3Exception):
    """
    Raised when message type isn't available. See i3.MSG_TYPES.
    """
    def __init__(self, type):
        msg = "Message type '%s' isn't available" % type
        super(MessageTypeError, self).__init__(msg)

class EventTypeError(i3Exception):
    """
    Raised when even type isn't available. See i3.EVENT_TYPES.
    """
    def __init__(self, type):
        msg = "Event type '%s' isn't available" % type
        super(EventTypeError, self).__init__(msg)

class MessageError(i3Exception):
    """
    Raised when a message to i3 is unsuccessful.
    That is, when it contains 'success': false in its JSON formatted response.
    """
    pass

class ConnectionError(i3Exception):
    """
    Raised when a socket couldn't connect to the window manager.
    """
    def __init__(self, socket_path):
        msg = "Could not connect to socket at '%s'" % socket_path
        super(ConnectionError, self).__init__(msg)


def parse_msg_type(msg_type):
    """
    Returns an i3-ipc code of the message type. Raises an exception if
    the given message type isn't available.
    """
    try:
        index = int(msg_type)
    except ValueError:
        index = -1
    if index >= 0 and index < len(MSG_TYPES):
        return index
    msg_type = str(msg_type).lower()
    if msg_type in MSG_TYPES:
        return MSG_TYPES.index(msg_type)
    else:
        raise MessageTypeError(msg_type)

def parse_event_type(event_type):
    """
    Returns an i3-ipc string of the event_type. Raises an exception if
    the given event type isn't available.
    """
    try:
        index = int(event_type)
    except ValueError:
        index = -1
    if index >= 0 and index < len(EVENT_TYPES):
        return EVENT_TYPES[index]
    event_type = str(event_type).lower()
    if event_type in EVENT_TYPES:
        return event_type
    else:
        raise EventTypeError(event_type)


class Socket(object):
    """
    Socket for communicating with the i3 window manager.
    Optional arguments:
    - 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
    timeout = 0.5  # in seconds
    buffer = b''  # byte string
    
    def __init__(self, path=None, timeout=None, chunk_size=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_size = struct.calcsize(self.struct_header)
    
    def initialize(self):
        """
        Initializes the socket.
        """
        self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.socket.settimeout(self.timeout)
    
    def connect(self, path=None):
        """
        Connects the socket to socket path if not already connected.
        """
        if not self.connected:
            self.initialize()
            if not path:
                path = self.path
            try:
                self.socket.connect(path)
            except socket.error:
                raise ConnectionError(path)
    
    def get(self, msg_type, payload=''):
        """
        Convenience method, calls "socket.send(msg_type, payload)" and
        returns data from "socket.receive()".
        """
        self.send(msg_type, payload)
        return self.receive()
    
    def subscribe(self, event_type, event=None):
        """
        Subscribes to an event. Returns data on first occurrence.
        """
        event_type = parse_event_type(event_type)
        # Create JSON payload from given event type and event
        payload = [event_type]
        if event:
            payload.append(event)
        payload = json.dumps(payload)
        return self.get('subscribe', payload)
    
    def send(self, msg_type, payload=''):
        """
        Sends the given message type with given message by packing them
        and continuously sending bytes from the packed message.
        """
        message = self.pack(msg_type, payload)
        # Continuously send the bytes from the message
        self.socket.sendall(message)
    
    def receive(self):
        """
        Tries to receive a data. Unpacks the received byte string if
        successful. Returns None on failure.
        """
        try:
            data = self.socket.recv(self.chunk_size)
            msg_magic, msg_length, msg_type = self.unpack_header(data)
            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.buffer + data
            return self.unpack(data)
        except socket.timeout:
            return None
    
    def pack(self, msg_type, payload):
        """
        Packs the given message type and payload. Turns the resulting
        message into a byte string.
        """
        msg_magic = self.magic_string
        # Get the byte count instead of number of characters
        msg_length = len(payload.encode('utf-8'))
        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)
        # Encoding the message back to byte string
        return message.encode('utf-8')
    
    def unpack(self, data):
        """
        Unpacks the given byte string and parses the result from JSON.
        Returns None on failure and saves data into "self.buffer".
        """
        data_size = len(data)
        msg_magic, msg_length, msg_type = self.unpack_header(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)
            self.buffer = data[msg_size:]
            return payload
        else:
            self.buffer = data
            return None
    
    def unpack_header(self, data):
        """
        Unpacks the header of given byte string.
        """
        return struct.unpack(self.struct_header, data[:self.struct_header_size])
    
    @property
    def connected(self):
        """
        Returns True if connected and False if not.
        """
        try:
            self.get('command')
            return True
        except socket.error:
            return False
    
    def close(self):
        """
        Closes the socket connection.
        """
        self.socket.close()


class Subscription(threading.Thread):
    """
    Creates a new subscription and runs a listener loop. Calls the
    callback on event.
    Example parameters:
    callback = lambda event, data, subscription: print(data)
    event_type = 'workspace'
    event = 'focus'
    event_socket = <i3.Socket object>
    data_socket = <i3.Socket object>
    """
    subscribed = False
    type_translation = {
        'workspace': 'get_workspaces',
        'output': 'get_outputs'
    }
    
    def __init__(self, callback, event_type, event=None, event_socket=None,
                 data_socket=None):
        # Variable initialization
        if not callable(callback):
            raise TypeError('Callback must be callable')
        event_type = parse_event_type(event_type)
        self.callback = callback
        self.event_type = event_type
        self.event = event
        # Socket initialization
        if not event_socket:
            event_socket = Socket()
        self.event_socket = event_socket
        self.event_socket.subscribe(event_type, event)
        if not data_socket:
            data_socket = Socket()
        self.data_socket = data_socket
        # Thread initialization
        threading.Thread.__init__(self)
        self.start()
    
    def run(self):
        """
        Wrapper method for the listen method -- handles exceptions.
        The method is run by the underlying "threading.Thread" object.
        """
        try:
            self.listen()
        except socket.error:
            self.close()
    
    def listen(self):
        """
        Runs a listener loop until self.subscribed is set to False.
        Calls the given callback method with data and the object itself.
        If event matches the given one, then matching data is retrieved.
        Otherwise, the event itself is sent to the callback.
        In that case 'change' key contains the thing that was changed.
        """
        self.subscribed = True
        while self.subscribed: