Allow matching tabs/windows over a custom all set

Author: kovidgoyal

kitty/boss.py

@@ -532,17 +532,19 @@ def all_windows(self) -> Iterator[Window]:
         for tab in self.all_tabs:
             yield from tab
 
-    def match_windows(self, match: str, self_window: Optional['Window'] = None) -> Iterator[Window]:
+    def match_windows(self, match: str, self_window: Optional['Window'] = None, all_windows: Iterable[Window] | None = None) -> Iterator[Window]:
+        all_windows = self.all_windows if all_windows is None else all_windows
         if match == 'all':
-            yield from self.all_windows
+            yield from all_windows
             return
         from .search_query_parser import search
         tab = self.active_tab
         if current_focused_os_window_id() <= 0:
             tm = self.os_window_map.get(last_focused_os_window_id())
             if tm is not None:
                 tab = tm.active_tab
-        window_id_limit = max(self.window_id_map, default=-1) + 1
+        wids = {w.id for w in all_windows}
+        window_id_limit = max(wids, default=-1) + 1
         active_session = self.active_session
 
         def get_matches(location: str, query: str, candidates: set[int]) -> set[int]:
@@ -557,25 +559,19 @@ def get_matches(location: str, query: str, candidates: set[int]) -> set[int]:
 
         for wid in search(match, (
             'id', 'title', 'pid', 'cwd', 'cmdline', 'num', 'env', 'var', 'recent', 'state', 'neighbor', 'session',
-        ), set(self.window_id_map), get_matches):
+        ), wids, get_matches):
             yield self.window_id_map[wid]
 
-    def tab_for_window(self, window: Window) -> Tab | None:
-        for tab in self.all_tabs:
-            for w in tab:
-                if w.id == window.id:
-                    return tab
-        return None
-
-    def match_tabs(self, match: str) -> Iterator[Tab]:
+    def match_tabs(self, match: str, all_tabs: Iterable[Tab] | None = None) -> Iterator[Tab]:
+        all_tabs = self.all_tabs if all_tabs is None else all_tabs
         if match == 'all':
-            yield from self.all_tabs
+            yield from all_tabs
             return
         from .search_query_parser import search
         tm = self.active_tab_manager
         if current_focused_os_window_id() <= 0:
             tm = self.os_window_map.get(last_focused_os_window_id()) or tm
-        tim = {t.id: t for t in self.all_tabs}
+        tim = {t.id: t for t in all_tabs}
         tab_id_limit = max(tim, default=-1) + 1
         window_id_limit = max(self.window_id_map, default=-1) + 1
 
@@ -592,13 +588,14 @@ def get_matches(location: str, query: str, candidates: set[int]) -> set[int]:
 
         found = False
         for tid in search(match, (
-                'id', 'index', 'title', 'window_id', 'window_title', 'pid', 'cwd', 'env', 'var', 'cmdline', 'recent', 'state', 'session',
+            'id', 'index', 'title', 'window_id', 'window_title', 'pid', 'cwd', 'env', 'var',
+            'cmdline', 'recent', 'state', 'session',
         ), set(tim), get_matches):
             found = True
             yield tim[tid]
 
         if not found:
-            tabs = {self.tab_for_window(w) for w in self.match_windows(match)}
+            tabs = {w.tabref() for w in self.match_windows(match)}
             for q in tabs:
                 if q:
                     yield q
@@ -3007,7 +3004,7 @@ def _move_window_to(
         window = window or self.active_window
         if not window:
             return
-        src_tab = self.tab_for_window(window)
+        src_tab = window.tabref()
         if src_tab is None:
             return
         with self.suppress_focus_change_events():

kitty/rc/base.py

@@ -390,8 +390,7 @@ def tabs_for_match_payload(self, boss: 'Boss', window: Optional['Window'], paylo
                 raise MatchError(match, 'tabs')
             return tabs
         if window and payload_get('self') in (None, True):
-            q = boss.tab_for_window(window)
-            if q:
+            if q := window.tabref():
                 return [q]
         t = boss.active_tab
         if t:

kitty/window.py

@@ -1467,8 +1467,7 @@ def add_kitten_result_processor(self, callback: Callable[['Window', Any], None])
         self.kitten_result_processors.append(callback)
 
     def handle_overlay_ready(self, msg: memoryview) -> None:
-        boss = get_boss()
-        tab = boss.tab_for_window(self)
+        tab = self.tabref()
         if tab is not None:
             tab.move_window_to_top_of_group(self)
         if self.keys_redirected_till_ready_from: