From ca03baf355324560537e12184087212cc24a4178 Mon Sep 17 00:00:00 2001
From: jwortmann <jwortmann@outlook.com>
Date: Tue, 12 Mar 2024 21:55:47 +0100
Subject: [PATCH] Keep tab selection and focus when applying WorkspaceEdit
 (#2431)

---
 plugin/core/sessions.py | 21 ++++++++++++++++-----
 plugin/hover.py         |  6 ++++--
 tests/setup.py          |  4 ++--
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/plugin/core/sessions.py b/plugin/core/sessions.py
index 1f4c08f74..827ba788c 100644
--- a/plugin/core/sessions.py
+++ b/plugin/core/sessions.py
@@ -1771,8 +1771,7 @@ def _apply_code_action_async(
             arguments = command.get("arguments")
             if arguments is not None:
                 execute_command['arguments'] = arguments
-            return promise.then(
-                lambda _: self.execute_command(execute_command, progress=False, view=view))
+            return promise.then(lambda _: self.execute_command(execute_command, progress=False, view=view))
         return promise
 
     def apply_workspace_edit_async(self, edit: WorkspaceEdit) -> Promise[None]:
@@ -1783,12 +1782,16 @@ def apply_workspace_edit_async(self, edit: WorkspaceEdit) -> Promise[None]:
         return self.apply_parsed_workspace_edits(parse_workspace_edit(edit))
 
     def apply_parsed_workspace_edits(self, changes: WorkspaceChanges) -> Promise[None]:
+        active_sheet = self.window.active_sheet()
+        selected_sheets = self.window.selected_sheets()
         promises = []  # type: List[Promise[None]]
         for uri, (edits, view_version) in changes.items():
             promises.append(
                 self.open_uri_async(uri).then(functools.partial(self._apply_text_edits, edits, view_version, uri))
             )
-        return Promise.all(promises).then(lambda _: None)
+        return Promise.all(promises) \
+            .then(lambda _: self._set_selected_sheets(selected_sheets)) \
+            .then(lambda _: self._set_focused_sheet(active_sheet))
 
     def _apply_text_edits(
         self, edits: List[TextEdit], view_version: Optional[int], uri: str, view: Optional[sublime.View]
@@ -1798,6 +1801,14 @@ def _apply_text_edits(
             return
         apply_text_edits(view, edits, required_view_version=view_version)
 
+    def _set_selected_sheets(self, sheets: List[sublime.Sheet]) -> None:
+        if len(sheets) > 1 and len(self.window.selected_sheets()) != len(sheets):
+            self.window.select_sheets(sheets)
+
+    def _set_focused_sheet(self, sheet: Optional[sublime.Sheet]) -> None:
+        if sheet and sheet != self.window.active_sheet():
+            self.window.focus_sheet(sheet)
+
     def decode_semantic_token(
             self, token_type_encoded: int, token_modifiers_encoded: int) -> Tuple[str, List[str], Optional[str]]:
         types_legend = tuple(cast(List[str], self.get_capability('semanticTokensProvider.legend.tokenTypes')))
@@ -1926,8 +1937,8 @@ def m_workspace_configuration(self, params: Dict[str, Any], request_id: Any) ->
 
     def m_workspace_applyEdit(self, params: Any, request_id: Any) -> None:
         """handles the workspace/applyEdit request"""
-        self.apply_workspace_edit_async(params.get('edit', {})).then(
-            lambda _: self.send_response(Response(request_id, {"applied": True})))
+        self.apply_workspace_edit_async(params.get('edit', {})) \
+            .then(lambda _: self.send_response(Response(request_id, {"applied": True})))
 
     def m_workspace_codeLens_refresh(self, _: Any, request_id: Any) -> None:
         """handles the workspace/codeLens/refresh request"""
diff --git a/plugin/hover.py b/plugin/hover.py
index d91975028..cde9540c6 100644
--- a/plugin/hover.py
+++ b/plugin/hover.py
@@ -195,8 +195,10 @@ def request_document_link_async(self, listener: AbstractViewListener, point: int
             if target:
                 link_promises.append(Promise.resolve(link))
             elif sv.has_capability_async("documentLinkProvider.resolveProvider"):
-                link_promises.append(sv.session.send_request_task(Request.resolveDocumentLink(link, sv.view)).then(
-                    lambda link: self._on_resolved_link(sv.session_buffer, link)))
+                link_promises.append(
+                    sv.session.send_request_task(Request.resolveDocumentLink(link, sv.view))
+                    .then(lambda link: self._on_resolved_link(sv.session_buffer, link))
+                )
         if link_promises:
             Promise.all(link_promises).then(partial(self._on_all_document_links_resolved, listener, point))
 
diff --git a/tests/setup.py b/tests/setup.py
index b72f5c134..059b4b7a6 100644
--- a/tests/setup.py
+++ b/tests/setup.py
@@ -210,8 +210,8 @@ def await_promise(cls, promise: Union[YieldPromise, Promise]) -> Generator:
     def await_run_code_action(self, code_action: Dict[str, Any]) -> Generator:
         promise = YieldPromise()
         sublime.set_timeout_async(
-            lambda: self.session.run_code_action_async(code_action, progress=False, view=self.view).then(
-                promise.fulfill))
+            lambda: self.session.run_code_action_async(code_action, progress=False, view=self.view)
+            .then(promise.fulfill))
         yield from self.await_promise(promise)
 
     def set_response(self, method: str, response: Any) -> None: