expose aio.Connection.subscribe

Author: Tony Crisci

i3ipc/aio/connection.py

@@ -247,7 +247,7 @@ def __init__(self, socket_path: Optional[str] = None, auto_reconnect: bool = Fal
         self._socket_path = socket_path
         self._auto_reconnect = auto_reconnect
         self._pubsub = _AIOPubSub(self)
-        self._subscriptions = 0
+        self._subscriptions = set()
         self._main_future = None
         self._reconnect_future = None
 
@@ -355,7 +355,7 @@ async def connect(self) -> 'Connection':
         self._sub_fd = self._sub_socket.fileno()
         self._loop.add_reader(self._sub_fd, self._message_reader)
 
-        await self._subscribe(self._subscriptions, force=True)
+        await self.subscribe(list(self._subscriptions), force=True)
 
         return self
 
@@ -419,25 +419,45 @@ async def _message(self, message_type: MessageType, payload: str = '') -> bytes:
 
         return message
 
-    async def _subscribe(self, events: Union[EventType, int], force=False):
+    async def subscribe(self, events: Union[List[Event], List[str]], force: bool = False):
+        """Send a ``SUBSCRIBE`` command to the ipc subscription connection and
+        await the result. To attach event handlers, use :func:`Connection.on()
+        <i3ipc.aio.Connection.on()>`. Calling this is only needed if you want
+        to be notified when events will start coming in.
+
+        :ivar events: A list of events to subscribe to. Currently you cannot
+            subscribe to detailed events.
+        :vartype events: list(:class:`Event <i3ipc.Event>`) or list(str)
+        :ivar force: If ``False``, the message will not be sent if this
+            connection is already subscribed to the event.
+        :vartype force: bool
+        """
         if not events:
             return
 
-        if type(events) is int:
-            events = EventType(events)
+        if type(events) is not list:
+            raise TypeError('events must be a list of events')
+
+        subscriptions = set()
+
+        for e in events:
+            e = Event(e)
+            if e not in Event._subscribable_events:
+                correct_event = str.split(e.value, '::')[0].upper()
+                raise ValueError(
+                    f'only nondetailed events are subscribable (use Event.{correct_event})')
+            subscriptions.add(e)
 
         if not force:
-            new_subscriptions = EventType(self._subscriptions ^ events.value)
-        else:
-            new_subscriptions = events
+            subscriptions = subscriptions.difference(self._subscriptions)
+            if not subscriptions:
+                return
 
-        if not new_subscriptions:
-            return
+        self._subscriptions.update(subscriptions)
+
+        payload = json.dumps([s.value for s in subscriptions])
 
-        self._subscriptions |= new_subscriptions.value
-        event_list = new_subscriptions.to_list()
-        await self._loop.sock_sendall(self._sub_socket,
-                                      _pack(MessageType.SUBSCRIBE, json.dumps(event_list)))
+        await self._loop.sock_sendall(self._sub_socket, _pack(MessageType.SUBSCRIBE, payload))
 
     def on(self, event: Union[Event, str], handler: Callable[['Connection', IpcBaseEvent], None]):
         """Subscribe to the event and call the handler when it is emitted by
@@ -456,9 +476,8 @@ def on(self, event: Union[Event, str], handler: Callable[['Connection', IpcBaseE
         if event.count('::') > 0:
             [event, __] = event.split('::')
 
-        event_type = EventType.from_string(event)
         self._pubsub.subscribe(event, handler)
-        asyncio.ensure_future(self._subscribe(event_type))
+        asyncio.ensure_future(self.subscribe([event]))
 
     def off(self, handler: Callable[['Connection', IpcBaseEvent], None]):
         """Unsubscribe the handler from being called on ipc events.

i3ipc/events.py

@@ -42,6 +42,12 @@ class Event(Enum):
     SHUTDOWN_EXIT = 'shutdown::exit'
 
 
+Event._subscribable_events = [
+    Event.WORKSPACE, Event.OUTPUT, Event.MODE, Event.WINDOW, Event.BARCONFIG_UPDATE, Event.BINDING,
+    Event.SHUTDOWN, Event.TICK
+]
+
+
 class WorkspaceEvent(IpcBaseEvent):
     """Sent when the user switches to a different workspace, when a new
     workspace is initialized or when a workspace is removed (because the last

test/aio/ipctest.py

@@ -8,7 +8,6 @@
 from random import random
 import asyncio
 
-from Xlib import display
 from .window import Window
 
 
@@ -64,11 +63,9 @@ async def command_checked(self, cmd):
 
         return result
 
-    async def open_window(self):
-        d = display.Display()
-        window = Window(d)
+    def open_window(self):
+        window = Window()
         window.run()
-        await self.i3_conn.command('nop')
         return window.window.id
 
     async def fresh_workspace(self):

test/aio/test_get_marks.py

@@ -6,10 +6,10 @@
 class TestGetMarks(IpcTest):
     @pytest.mark.asyncio
     async def test_get_marks(self, i3):
-        await self.open_window()
+        self.open_window()
         await self.command_checked('mark a')
         await self.command_checked('mark --add b')
-        await self.open_window()
+        self.open_window()
         await self.command_checked('mark "(╯°□°）╯︵ ┻━┻"')
 
         marks = await i3.get_marks()

test/aio/test_leaves.py

@@ -7,10 +7,10 @@ class TestLeaves(IpcTest):
     @pytest.mark.asyncio
     async def test_workspace_leaves(self, i3):
         ws_name = await self.fresh_workspace()
-        con1 = await self.open_window()
+        con1 = self.open_window()
         await self.command_checked(f'[id={con1}] floating enable')
-        await self.open_window()
-        await self.open_window()
+        self.open_window()
+        self.open_window()
 
         tree = await i3.get_tree()
         ws = [w for w in tree.workspaces() if w.name == ws_name][0]

test/aio/test_window.py

@@ -1,7 +1,7 @@
 from .ipctest import IpcTest
 
 import pytest
-import asyncio
+from i3ipc import Event
 
 
 class TestWindow(IpcTest):
@@ -14,12 +14,10 @@ def on_window(self, i3, e):
     @pytest.mark.asyncio
     async def test_window_event(self, i3):
         await self.fresh_workspace()
-        i3.on('window', self.on_window)
+        await i3.subscribe([Event.WINDOW])
+        i3.on(Event.WINDOW, self.on_window)
 
-        def open_window():
-            asyncio.ensure_future(self.open_window())
-
-        i3._loop.call_later(0.1, open_window)
+        self.open_window()
 
         await i3.main()
 
@@ -28,7 +26,7 @@ def open_window():
     @pytest.mark.asyncio
     async def test_marks(self, i3):
         await self.fresh_workspace()
-        await self.open_window()
+        self.open_window()
         await i3.command('mark foo')
         tree = await i3.get_tree()
         assert 'foo' in tree.find_focused().marks