From 85fe1f9889ba34fdb43d08b264c65abd3a8693be Mon Sep 17 00:00:00 2001
From: rysson <robert.kalinowski@sharkbits.com>
Date: Tue, 21 Jun 2022 16:16:24 +0200
Subject: [PATCH] Support for restart i3 reconnection.

Search the socket path on reconnect if no user's socket_path defined.
---
 i3ipc/aio/connection.py | 10 +++++++---
 i3ipc/connection.py     | 37 +++++++++++++++++++++----------------
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/i3ipc/aio/connection.py b/i3ipc/aio/connection.py
index c4e65e8..66f587d 100644
--- a/i3ipc/aio/connection.py
+++ b/i3ipc/aio/connection.py
@@ -277,6 +277,7 @@ class Connection:
     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()
@@ -297,7 +298,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:
@@ -393,8 +394,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(self._display)
+        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')
@@ -428,6 +431,7 @@ async def do_reconnect():
                     error = None
                     break
                 except Exception as e:
+                    self._connection_socket_path = None
                     error = e
                     await asyncio.sleep(0.001)
 
diff --git a/i3ipc/connection.py b/i3ipc/connection.py
index 778e8ba..c91df5b 100644
--- a/i3ipc/connection.py
+++ b/i3ipc/connection.py
@@ -59,24 +59,30 @@ 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')
 
         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
-        self._display = display
+
+    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")
@@ -134,7 +140,7 @@ def socket_path(self) -> str:
 
         :rtype: str
         """
-        return self._socket_path
+        return self._connection_socket_path
 
     @property
     def auto_reconnect(self) -> bool:
@@ -195,14 +201,13 @@ 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
+            self._connection_socket_path = self._get_socket_path()
+            if os.path.exists(self._connection_socket_path):
+                return True
             time.sleep(0.001)
 
-        return socket_path_exists
+        return False
 
     def _message(self, message_type, payload):
         try:
@@ -213,13 +218,13 @@ def _message(self, message_type, payload):
                 raise e
 
             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()
@@ -475,7 +480,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)