A custom DISPLAY to determinate the socket_path.

i3ipc/aio/connection.py

@@ -177,7 +177,7 @@ def _unpack_header(data: bytes) -> Tuple[bytes, int, int]:
     return struct.unpack(_struct_header, data[:_struct_header_size])
 
 
-async def _find_socket_path() -> Optional[str]:
+async def _find_socket_path(disp: Optional[str] = None) -> Optional[str]:
     socket_path = None
 
     def exists(path):
@@ -202,12 +202,17 @@ def exists(path):
             return socket_path
 
     # finally try the binaries
+    if disp:
+        env = {**os.environ, 'DISPLAY': disp}
+    else:
+        env = None
     for binary in ('i3', 'sway'):
         try:
             process = await asyncio.create_subprocess_exec(binary,
                                                            '--get-socketpath',
                                                            stdout=PIPE,
-                                                           stderr=PIPE)
+                                                           stderr=PIPE,
+                                                           env=env)
 
             stdout, stderr = await process.communicate()
 
@@ -251,17 +256,22 @@ class Connection:
     :param auto_reconnect: Whether to attempt to reconnect if the connection to
         the socket is broken when i3 restarts.
     :type auto_reconnect: bool
+    :param display: A custom DISPLAY to determinate the socket_path.
+    :type display: str
 
     :raises Exception: If the connection to i3 cannot be established.
     """
-    def __init__(self, socket_path: Optional[str] = None, auto_reconnect: bool = False):
+    def __init__(self, socket_path: Optional[str] = None, auto_reconnect: bool = False, *,
+                 display: Optional[str] = None):
         self._socket_path = socket_path
+        self._connection_socket_path = None
         self._auto_reconnect = auto_reconnect
         self._pubsub = _AIOPubSub(self)
         self._subscriptions = set()
         self._main_future = None
         self._reconnect_future = None
         self._synchronizer = None
+        self._display = display
 
     def _sync(self):
         if self._synchronizer is None:
@@ -275,7 +285,7 @@ def socket_path(self) -> str:
 
         :rtype: str
         """
-        return self._socket_path
+        return self._connection_socket_path or self._socket_path
 
     @property
     def auto_reconect(self) -> bool:
@@ -371,8 +381,10 @@ async def connect(self) -> 'Connection':
         if self._socket_path:
             logger.info('using user provided socket path: {}', self._socket_path)
 
-        if not self._socket_path:
-            self._socket_path = await _find_socket_path()
+        if self._socket_path:
+            self._connection_socket_path = self._socket_path
+        else:
+            self._connection_socket_path = await _find_socket_path(self._display)
 
         if not self.socket_path:
             raise Exception('Failed to retrieve the i3 or sway IPC socket path')
@@ -406,6 +418,7 @@ async def do_reconnect():
                     error = None
                     break
                 except Exception as e:
+                    self._connection_socket_path = None
                     error = e
                     await asyncio.sleep(0.001)
 

i3ipc/connection.py

@@ -41,6 +41,8 @@ class Connection:
     :param auto_reconnect: Whether to attempt to reconnect if the connection to
         the socket is broken when i3 restarts.
     :type auto_reconnect: bool
+    :param display: A custom DISPLAY to determinate the socket_path.
+    :type display: str
 
     :raises Exception: If the connection to i3 cannot be established.
     """
@@ -50,28 +52,36 @@ class Connection:
     _struct_header = '=%dsII' % len(_MAGIC.encode('utf-8'))
     _struct_header_size = struct.calcsize(_struct_header)
 
-    def __init__(self, socket_path=None, auto_reconnect=False):
+    def __init__(self, socket_path=None, auto_reconnect=False, display=None):
 
         if socket_path:
             logger.info('using user provided socket path: %s', socket_path)
-        else:
-            socket_path = self._find_socket_path()
-
-        if not socket_path:
-            raise Exception('Failed to retrieve the i3 or sway IPC socket path')
+        # else _find_socket_path() will be used
 
         self.subscriptions = 0
         self._pubsub = PubSub(self)
+        self._display = display
         self._socket_path = socket_path
+        self._connection_socket_path = self._get_socket_path()
         self._cmd_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        self._cmd_socket.connect(self._socket_path)
+        self._cmd_socket.connect(self._connection_socket_path)
         self._cmd_lock = Lock()
         self._sub_socket = None
         self._sub_lock = Lock()
         self._auto_reconnect = auto_reconnect
         self._quitting = False
         self._synchronizer = None
 
+    def _get_socket_path(self):
+        """Returns a current socket path."""
+        if self._socket_path:
+            socket_path = self._socket_path
+        else:
+            socket_path = self._find_socket_path()
+        if not socket_path:
+            raise Exception('Failed to retrieve the i3 or sway IPC socket path')
+        return socket_path
+
     def _find_socket_path(self):
         socket_path = os.environ.get("I3SOCK")
         if socket_path:
@@ -83,9 +93,13 @@ def _find_socket_path(self):
             logger.info('got socket path from SWAYSOCK env variable: %s', socket_path)
             return socket_path
 
+        if self._display:
+            env = {**os.environ, 'DISPLAY': self._display}
+        else:
+            env = None
         for binary in ('i3', 'sway'):
             try:
-                process = run([binary, '--get-socketpath'], stdout=PIPE, stderr=PIPE)
+                process = run([binary, '--get-socketpath'], stdout=PIPE, stderr=PIPE, env=env)
                 if process.returncode == 0 and process.stdout:
                     socket_path = process.stdout.decode().strip()
                     logger.info('got socket path from `%s` binary: %s', binary, socket_path)
@@ -114,7 +128,7 @@ def socket_path(self) -> str:
 
         :rtype: str
         """
-        return self._socket_path
+        return self._connection_socket_path or self._socket_path
 
     @property
     def auto_reconnect(self) -> bool:
@@ -175,14 +189,16 @@ def _ipc_send(self, sock, message_type, payload):
 
     def _wait_for_socket(self):
         # for the auto_reconnect feature only
-        socket_path_exists = False
         for tries in range(0, 500):
-            socket_path_exists = os.path.exists(self._socket_path)
-            if socket_path_exists:
-                break
+            try:
+                self._connection_socket_path = self._get_socket_path()
+                if os.path.exists(self._connection_socket_path):
+                    return True
+            except Exception:
+                pass
             time.sleep(0.001)
 
-        return socket_path_exists
+        return False
 
     def _message(self, message_type, payload):
         try:
@@ -191,15 +207,16 @@ def _message(self, message_type, payload):
         except ConnectionError as e:
             if not self.auto_reconnect:
                 raise e
+            self._connection_socket_path = None
 
             logger.info('got a connection error, reconnecting', exc_info=e)
-            # XXX: can the socket path change between restarts?
+            # The socket path can change between restarts.
             if not self._wait_for_socket():
                 logger.info('could not reconnect')
                 raise e
 
             self._cmd_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-            self._cmd_socket.connect(self._socket_path)
+            self._cmd_socket.connect(self._connection_socket_path)
             return self._ipc_send(self._cmd_socket, message_type, payload)
         finally:
             self._cmd_lock.release()
@@ -455,7 +472,7 @@ def _on(self, event: Union[Event, str], handler: Callable[['Connection', IpcBase
 
     def _event_socket_setup(self):
         self._sub_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-        self._sub_socket.connect(self._socket_path)
+        self._sub_socket.connect(self._connection_socket_path)
 
         self._subscribe(self.subscriptions)