Implement full serialization/unserialization for layouts

Need to go over all the non-split layouts and see if they need
any TLC

kitty/constants.py

@@ -34,6 +34,7 @@ default_pager_for_help = ('less', '-iRXF')
 kitty_run_data: dict[str, Any] = getattr(sys, 'kitty_run_data', {})
 launched_by_launch_services = kitty_run_data.get('launched_by_launch_services', False)
 is_quick_access_terminal_app = kitty_run_data.get('is_quick_access_terminal_app', False)
+serialize_user_var_name = 'kitty_serialize_window_id'
 
 if getattr(sys, 'frozen', False):
     extensions_dir: str = kitty_run_data['extensions_dir']

kitty/layout/base.py

@@ -7,10 +7,11 @@ from itertools import repeat
 from typing import Any, NamedTuple
 
 from kitty.borders import BorderColor
+from kitty.constants import serialize_user_var_name
 from kitty.fast_data_types import Region, set_active_window, viewport_for_window
 from kitty.options.types import Options
-from kitty.types import Edges, WindowGeometry
-from kitty.typing_compat import TypedDict, WindowMapper, WindowType
+from kitty.types import Edges, WindowGeometry, WindowMapper
+from kitty.typing_compat import TypedDict, WindowType
 from kitty.window_list import WindowGroup, WindowList
 
 
@@ -437,16 +438,33 @@ class Layout:
     def layout_state(self) -> dict[str, Any]:
         return {}
 
-    def set_layout_state(self, layout_state: dict[str, Any], map_window_id: WindowMapper) -> bool:
+    def set_layout_state(self, layout_state: dict[str, Any], map_group_id: WindowMapper) -> bool:
         return True
 
-    def serialize(self) -> dict[str, Any]:
+    def serialize(self, all_windows: WindowList) -> dict[str, Any]:
         ans = self.layout_state()
         ans['opts'] = self.layout_opts.serialized()
         ans['class'] = self.__class__.__name__
+        ans['all_windows'] = aw = all_windows.serialize_state()
+        for wg in aw['window_groups']:
+            wg['window_ids'] = tuple(w['id'] for w in aw.pop('windows'))
         return ans
 
-    def unserialize(self, s: dict[str, Any], map_window_id: WindowMapper) -> bool:
+    def unserialize(
+        self, s: dict[str, Any], all_windows: WindowList,
+        serialize_user_var_name: str = serialize_user_var_name
+    ) -> bool:
         if s.get('class') != self.__class__.__name__:
             return False
-        return self.set_layout_state(s, map_window_id)
+        window_id_map = {}
+        for w in all_windows:
+            k = w.user_vars.pop(serialize_user_var_name, None)
+            if k is not None:
+                try:
+                    window_id_map[int(k)] = w.id
+                except Exception:
+                    pass
+        m = all_windows.unserialize_layout_state(s['all_windows'], window_id_map)
+        if m is None:
+            return False
+        return self.set_layout_state(s, m.get)

kitty/layout/splits.py

@@ -5,8 +5,8 @@ from collections.abc import Collection, Generator, Sequence
 from typing import Any, NamedTuple, Optional, TypedDict, Union
 
 from kitty.borders import BorderColor
-from kitty.types import Edges, WindowGeometry
-from kitty.typing_compat import EdgeLiteral, WindowMapper, WindowType
+from kitty.types import Edges, WindowGeometry, WindowMapper
+from kitty.typing_compat import EdgeLiteral, WindowType
 from kitty.window_list import WindowGroup, WindowList
 
 from .base import BorderLine, Layout, LayoutOpts, NeighborsMap, blank_rects_for_window, lgd, window_geometry_from_layouts
@@ -55,8 +55,7 @@ class Pair:
             if x is None:
                 return None
             if isinstance(x, int):
-                w = map_window_id(x)
-                return None if w is None else w.id
+                return map_window_id(x)
             ans = Pair()
             ans.unserialize(x, map_window_id)
             return ans if ans.one or ans.two else None
@@ -710,9 +709,9 @@ class Splits(Layout):
     def layout_state(self) -> dict[str, Any]:
         return {'pairs': self.pairs_root.serialize()}
 
-    def set_layout_state(self, layout_state: dict[str, Any], map_window_id: WindowMapper) -> bool:
+    def set_layout_state(self, layout_state: dict[str, Any], map_group_id: WindowMapper) -> bool:
         new_root = Pair()
-        new_root.unserialize(layout_state['pairs'], map_window_id)
+        new_root.unserialize(layout_state['pairs'], map_group_id)
         before = frozenset(self.pairs_root.all_window_ids())
         if before == frozenset(new_root.all_window_ids()):
             self.pairs_root = new_root

kitty/tabs.py

@@ -272,7 +272,7 @@ class Tab:  # {{{
             'window_list': self.windows.serialize_state(),
             'current_layout': self._current_layout_name,
             'last_used_layout': self._last_used_layout,
-            'layout_state': self.current_layout.serialize(),
+            'layout_state': self.current_layout.serialize(self.windows),
             'enabled_layouts': self.enabled_layouts,
             'name': self.name,
         }

kitty/types.py

@@ -230,4 +230,5 @@ def ac(group: ActionGroup, doc: str) -> Callable[[_T], _T]:
     return w
 
 
+WindowMapper = Callable[[int], int | None]
 DecoratedFunc = TypeVar('DecoratedFunc', bound=Callable[..., Any])

kitty/typing_compat.py

@@ -23,4 +23,3 @@ Protocol = object
 OptionsProtocol = object
 NotRequired = object
 CoreTextFont = FontConfigPattern = dict
-WindowMapper = object

kitty/typing_compat.pyi

@@ -54,7 +54,6 @@ GRT_C = Literal[0, 1]
 GRT_d = Literal['a', 'A', 'c', 'C', 'i', 'I', 'p', 'P', 'q', 'Q', 'x', 'X', 'y', 'Y', 'z', 'Z', 'f', 'F']
 ReadableBuffer = bytes | bytearray | memoryview | array.array[int] | mmap.mmap
 WriteableBuffer = bytearray | memoryview | array.array[int] | mmap.mmap
-WindowMapper = Callable[[int], WindowType | None]
 
 
 
@@ -72,5 +71,5 @@ __all__ = (
     'KeyActionType', 'KeyMap', 'KittyCommonOpts', 'AliasMap', 'CoreTextFont', 'WindowSystemMouseEvent',
     'FontConfigPattern', 'ScreenType', 'StartupCtx', 'KeyEventType', 'LayoutType', 'PowerlineStyle',
     'RemoteCommandType', 'SessionType', 'SessionTab', 'SpecialWindowInstance', 'TabType', 'ScreenSize', 'WindowType',
-    'ReadableBuffer', 'WriteableBuffer', 'WindowMapper',
+    'ReadableBuffer', 'WriteableBuffer',
 )

kitty/window_list.py

@@ -6,7 +6,7 @@ from collections import deque
 from collections.abc import Iterator
 from contextlib import suppress
 from itertools import count
-from typing import Any, Deque, Union
+from typing import Any, Deque, Sequence, Union
 
 from .fast_data_types import Color, get_options
 from .types import OverlayType, WindowGeometry
@@ -48,6 +48,12 @@ class WindowGroup:
                 return True
         return False
 
+    def has_window_id(self, wid: int) -> bool:
+        for w in self.windows:
+            if w.id == wid:
+                return True
+        return False
+
     @property
     def needs_attention(self) -> bool:
         for w in self.windows:
@@ -93,6 +99,12 @@ class WindowGroup:
             'windows': [w.serialize_state() for w in self.windows]
         }
 
+    def unserialize_layout_state(self, window_ids: Sequence[int]) -> None:
+        order_map = {wid: i for i, wid in enumerate(window_ids)}
+        def sort_key(w: WindowType) -> int:
+            return order_map.get(w.id, -1)
+        self.windows.sort(key=sort_key)
+
     def as_simple_dict(self) -> dict[str, Any]:
         return {
             'id': self.id,
@@ -173,6 +185,55 @@ class WindowList:
             'window_groups': [g.serialize_state() for g in self.groups]
         }
 
+    def unserialize_layout_state(self, state: dict[str, Any], window_id_map: dict[int, int]) -> dict[int, int] | None:
+        if set(window_id_map.values()) != set(self.id_map):
+            # some window in this collection does not correspond to a
+            # serialized window
+            return None
+        ans = {}
+        gmap = {g.id: g for g in self.groups}
+        present_wids_map = {g.id: {w.id for w in g} for g in self.groups}
+
+        def unmapped_group_having_subset_of_windows(wids: Sequence[int]) -> Iterator[WindowGroup]:
+            mapped_wids = set()
+            for wid in wids:
+                new_wid = window_id_map.get(wid)
+                if new_wid is not None:
+                    mapped_wids.add(new_wid)
+            for gid in tuple(gmap):
+                present_wids = present_wids_map[gid]
+                if present_wids.issubset(mapped_wids):
+                    yield gmap.pop(gid)
+                    break
+
+        for wg in state['window_groups']:
+            old_group_id = wg['id']
+            for group in unmapped_group_having_subset_of_windows(wg['window_ids']):
+                ans[old_group_id] = group.id
+        # check that all the groups present were also in the serialized state.
+        # there could have been extra windows/groups in the serialized state,
+        # we ignore them.
+        if len(ans) != len(self.groups):
+            return None
+        gmap = {g.id: g for g in self.groups}
+        groups = []
+        for wg in state['window_groups']:
+            old_group_id = wg['id']
+            if new_group_id := ans.get(old_group_id):
+                groups.append((g := gmap[new_group_id]))
+                new_window_ids = []
+                for old_window_id in wg['window_ids']:
+                    if new_window_id := window_id_map.get(old_window_id):
+                        new_window_ids.append(new_window_id)
+                g.unserialize_layout_state(new_window_ids)
+        self.groups = groups
+        history = []
+        for old_wid in state['active_group_history']:
+            if new_wid := window_id_map.get(old_wid):
+                history.append(new_wid)
+        self.active_group_history = deque(history, 64)
+        return ans
+
     @property
     def active_group_idx(self) -> int:
         return self._active_group_idx