Preserve state while reloading directory tree.

Author: rodrigogiraoserrao

CHANGELOG.md

@@ -11,6 +11,14 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 - Fixed `DirectoryTree.clear_node` not clearing the node specified https://github.com/Textualize/textual/issues/4122
 
+### Added
+
+- `Tree` (and `DirectoryTree`) grew an attribute `lock` that can be used for synchronization across coroutines https://github.com/Textualize/textual/issues/4056
+
+### Changed
+
+- `DirectoryTree.reload` and `DirectoryTree.reload_node` now preserve state when reloading https://github.com/Textualize/textual/issues/4056
+
 ## [0.48.2] - 2024-02-02
 
 ### Fixed

src/textual/widgets/_directory_tree.py

@@ -5,20 +5,19 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, ClassVar, Iterable, Iterator
 
-from ..await_complete import AwaitComplete
-
-if TYPE_CHECKING:
-    from typing_extensions import Self
-
 from rich.style import Style
 from rich.text import Text, TextType
 
 from .. import work
+from ..await_complete import AwaitComplete
 from ..message import Message
 from ..reactive import var
 from ..worker import Worker, WorkerCancelled, WorkerFailed, get_current_worker
 from ._tree import TOGGLE_STYLE, Tree, TreeNode
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 
 @dataclass
 class DirEntry:
@@ -164,7 +163,7 @@ def _add_to_load_queue(self, node: TreeNode[DirEntry]) -> AwaitComplete:
 
         Returns:
             An optionally awaitable object that can be awaited until the
-            load queue has finished processing.
+                load queue has finished processing.
         """
         assert node.data is not None
         if not node.data.loaded:
@@ -174,16 +173,18 @@ def _add_to_load_queue(self, node: TreeNode[DirEntry]) -> AwaitComplete:
         return AwaitComplete(self._load_queue.join())
 
     def reload(self) -> AwaitComplete:
-        """Reload the `DirectoryTree` contents."""
-        self.reset(str(self.path), DirEntry(self.PATH(self.path)))
+        """Reload the `DirectoryTree` contents.
+
+        Returns:
+            An optionally awaitable that ensures the tree has finished reloading.
+        """
         # Orphan the old queue...
         self._load_queue = Queue()
+        # ... reset the root node ...
+        processed = self.reload_node(self.root)
         # ...and replace the old load with a new one.
         self._loader()
-        # We have a fresh queue, we have a fresh loader, get the fresh root
-        # loading up.
-        queue_processed = self._add_to_load_queue(self.root)
-        return queue_processed
+        return processed
 
     def clear_node(self, node: TreeNode[DirEntry]) -> Self:
         """Clear all nodes under the given node.
@@ -215,6 +216,88 @@ def reset_node(
         node.data = data
         return self
 
+    async def _reload(self, node: TreeNode[DirEntry]) -> None:
+        """Reloads the subtree rooted at the given node while preserving state.
+
+        After reloading the subtree, nodes that were expanded and still exist
+        will remain expanded and the highlighted node will be preserved, if it
+        still exists. If it doesn't, highlighting goes up to the first parent
+        directory that still exists.
+
+        Args:
+            node: The root of the subtree to reload.
+        """
+        async with self.lock:
+            # Track nodes that were expanded before reloading.
+            currently_open: set[Path] = set()
+            to_check: list[TreeNode[DirEntry]] = [node]
+            while to_check:
+                checking = to_check.pop()
+                if checking.allow_expand and checking.is_expanded:
+                    if checking.data:
+                        currently_open.add(checking.data.path)
+                    to_check.extend(checking.children)
+
+            # Track node that was highlighted before reloading.
+            highlighted_path: None | Path = None
+            if self.cursor_line > -1:
+                highlighted_node = self.get_node_at_line(self.cursor_line)
+                if highlighted_node is not None and highlighted_node.data is not None:
+                    highlighted_path = highlighted_node.data.path
+
+            if node.data is not None:
+                self.reset_node(
+                    node, str(node.data.path.name), DirEntry(self.PATH(node.data.path))
+                )
+
+            # Reopen nodes that were expanded and still exist.
+            to_reopen = [node]
+            while to_reopen:
+                reopening = to_reopen.pop()
+                if not reopening.data:
+                    continue
+                if (
+                    reopening.allow_expand
+                    and (reopening.data.path in currently_open or reopening == node)
+                    and reopening.data.path.exists()
+                ):
+                    try:
+                        content = await self._load_directory(reopening).wait()
+                    except (WorkerCancelled, WorkerFailed):
+                        continue
+                    reopening.data.loaded = True
+                    self._populate_node(reopening, content)
+                    to_reopen.extend(reopening.children)
+                    reopening.expand()
+
+            if highlighted_path is None:
+                return
+
+            # Restore the highlighted path and consider the parents as fallbacks.
+            looking = [node]
+            highlight_candidates = set(highlighted_path.parents)
+            highlight_candidates.add(highlighted_path)
+            best_found: None | TreeNode[DirEntry] = None
+            while looking:
+                checking = looking.pop()
+                checking_path = (
+                    checking.data.path if checking.data is not None else None
+                )
+                if checking_path in highlight_candidates:
+                    best_found = checking
+                    if checking_path == highlighted_path:
+                        break
+                if (
+                    checking.allow_expand
+                    and checking.is_expanded
+                    and checking_path in highlighted_path.parents
+                ):
+                    looking.extend(checking.children)
+            if best_found is not None:
+                # We need valid lines. Make sure the tree lines have been computed:
+                _ = self._tree_lines
+                self.cursor_line = best_found.line
+
     def reload_node(self, node: TreeNode[DirEntry]) -> AwaitComplete:
         """Reload the given node's contents.
 
@@ -223,12 +306,12 @@ def reload_node(self, node: TreeNode[DirEntry]) -> AwaitComplete:
         or any other nodes).
 
         Args:
-            node: The node to reload.
+            node: The root of the subtree to reload.
+
+        Returns:
+            An optionally awaitable that ensures the subtree has finished reloading.
         """
-        self.reset_node(
-            node, str(node.data.path.name), DirEntry(self.PATH(node.data.path))
-        )
-        return self._add_to_load_queue(node)
+        return AwaitComplete(self._reload(node))
 
     def validate_path(self, path: str | Path) -> Path:
         """Ensure that the path is of the `Path` type.
@@ -415,28 +498,29 @@ async def _loader(self) -> None:
             # this blocks if the queue is empty.
             node = await self._load_queue.get()
             content: list[Path] = []
-            try:
-                # Spin up a short-lived thread that will load the content of
-                # the directory associated with that node.
-                content = await self._load_directory(node).wait()
-            except WorkerCancelled:
-                # The worker was cancelled, that would suggest we're all
-                # done here and we should get out of the loader in general.
-                break
-            except WorkerFailed:
-                # This particular worker failed to start. We don't know the
-                # reason so let's no-op that (for now anyway).
-                pass
-            else:
-                # We're still here and we have directory content, get it into
-                # the tree.
-                if content:
-                    self._populate_node(node, content)
-            finally:
-                # Mark this iteration as done.
-                self._load_queue.task_done()
-
-    async def _on_tree_node_expanded(self, event: Tree.NodeExpanded) -> None:
+            async with self.lock:
+                try:
+                    # Spin up a short-lived thread that will load the content of
+                    # the directory associated with that node.
+                    content = await self._load_directory(node).wait()
+                except WorkerCancelled:
+                    # The worker was cancelled, that would suggest we're all
+                    # done here and we should get out of the loader in general.
+                    break
+                except WorkerFailed:
+                    # This particular worker failed to start. We don't know the
+                    # reason so let's no-op that (for now anyway).
+                    pass
+                else:
+                    # We're still here and we have directory content, get it into
+                    # the tree.
+                    if content:
+                        self._populate_node(node, content)
+                finally:
+                    # Mark this iteration as done.
+                    self._load_queue.task_done()
+
+    async def _on_tree_node_expanded(self, event: Tree.NodeExpanded[DirEntry]) -> None:
         event.stop()
         dir_entry = event.node.data
         if dir_entry is None:
@@ -446,7 +530,7 @@ async def _on_tree_node_expanded(self, event: Tree.NodeExpanded) -> None:
         else:
             self.post_message(self.FileSelected(event.node, dir_entry.path))
 
-    def _on_tree_node_selected(self, event: Tree.NodeSelected) -> None:
+    def _on_tree_node_selected(self, event: Tree.NodeSelected[DirEntry]) -> None:
         event.stop()
         dir_entry = event.node.data
         if dir_entry is None:

src/textual/widgets/_tree.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from asyncio import Lock
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar, Generic, Iterable, NewType, TypeVar, cast
 
@@ -615,8 +616,10 @@ def __init__(
         self.root = self._add_node(None, text_label, data)
         """The root node of the tree."""
         self._line_cache: LRUCache[LineCacheKey, Strip] = LRUCache(1024)
-        self._tree_lines_cached: list[_TreeLine] | None = None
+        self._tree_lines_cached: list[_TreeLine[TreeDataType]] | None = None
         self._cursor_node: TreeNode[TreeDataType] | None = None
+        self.lock = Lock()
+        """Used to synchronise stateful directory tree operations."""
 
         super().__init__(name=name, id=id, classes=classes, disabled=disabled)
 
@@ -815,7 +818,7 @@ def _invalidate(self) -> None:
         self.root._reset()
         self.refresh(layout=True)
 
-    def _on_mouse_move(self, event: events.MouseMove):
+    def _on_mouse_move(self, event: events.MouseMove) -> None:
         meta = event.style.meta
         if meta and "line" in meta:
             self.hover_line = meta["line"]
@@ -948,7 +951,7 @@ def _refresh_node(self, node: TreeNode[TreeDataType]) -> None:
                 self._refresh_line(line_no)
 
     @property
-    def _tree_lines(self) -> list[_TreeLine]:
+    def _tree_lines(self) -> list[_TreeLine[TreeDataType]]:
         if self._tree_lines_cached is None:
             self._build()
         assert self._tree_lines_cached is not None
@@ -957,13 +960,14 @@ def _tree_lines(self) -> list[_TreeLine]:
     async def _on_idle(self, event: events.Idle) -> None:
         """Check tree needs a rebuild on idle."""
         # Property calls build if required
-        self._tree_lines
+        async with self.lock:
+            self._tree_lines
 
     def _build(self) -> None:
         """Builds the tree by traversing nodes, and creating tree lines."""
 
         TreeLine = _TreeLine
-        lines: list[_TreeLine] = []
+        lines: list[_TreeLine[TreeDataType]] = []
         add_line = lines.append
 
         root = self.root
@@ -989,7 +993,7 @@ def add_node(
         show_root = self.show_root
         get_label_width = self.get_label_width
 
-        def get_line_width(line: _TreeLine) -> int:
+        def get_line_width(line: _TreeLine[TreeDataType]) -> int:
             return get_label_width(line.node) + line._get_guide_width(
                 guide_depth, show_root
             )
@@ -1147,17 +1151,18 @@ def _toggle_node(self, node: TreeNode[TreeDataType]) -> None:
             node.expand()
 
     async def _on_click(self, event: events.Click) -> None:
-        meta = event.style.meta
-        if "line" in meta:
-            cursor_line = meta["line"]
-            if meta.get("toggle", False):
-                node = self.get_node_at_line(cursor_line)
-                if node is not None:
-                    self._toggle_node(node)
+        async with self.lock:
+            meta = event.style.meta
+            if "line" in meta:
+                cursor_line = meta["line"]
+                if meta.get("toggle", False):
+                    node = self.get_node_at_line(cursor_line)
+                    if node is not None:
+                        self._toggle_node(node)
 
-            else:
-                self.cursor_line = cursor_line
-                await self.run_action("select_cursor")
+                else:
+                    self.cursor_line = cursor_line
+                    await self.run_action("select_cursor")
 
     def notify_style_update(self) -> None:
         self._invalidate()