From abca7bd8e074a57318f9f51e3b873abd6fc135af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rodrigo=20Gir=C3=A3o=20Serr=C3=A3o?=
<5621605+rodrigogiraoserrao@users.noreply.github.com>
Date: Mon, 5 Feb 2024 16:59:27 +0000
Subject: [PATCH] Preserve state while reloading directory tree.
---
CHANGELOG.md | 8 +
src/textual/widgets/_directory_tree.py | 164 +++++++++++++-----
src/textual/widgets/_tree.py | 37 ++--
.../__snapshots__/test_snapshots.ambr | 160 +++++++++++++++++
.../snapshot_apps/directory_tree_reload.py | 62 +++++++
tests/snapshot_tests/test_snapshots.py | 13 ++
6 files changed, 388 insertions(+), 56 deletions(-)
create mode 100644 tests/snapshot_tests/snapshot_apps/directory_tree_reload.py
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bb8c1a71bf1..acd43d635b4 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/src/textual/widgets/_directory_tree.py b/src/textual/widgets/_directory_tree.py
index 3cee5fca5a0..46bef1550ca 100644
--- a/src/textual/widgets/_directory_tree.py
+++ b/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:
diff --git a/src/textual/widgets/_tree.py b/src/textual/widgets/_tree.py
index 43a079a5ebf..3e0a2d8cb4c 100644
--- a/src/textual/widgets/_tree.py
+++ b/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()
diff --git a/tests/snapshot_tests/__snapshots__/test_snapshots.ambr b/tests/snapshot_tests/__snapshots__/test_snapshots.ambr
index a9d82110a11..9dec40dc683 100644
--- a/tests/snapshot_tests/__snapshots__/test_snapshots.ambr
+++ b/tests/snapshot_tests/__snapshots__/test_snapshots.ambr
@@ -16740,6 +16740,166 @@
'''
# ---
+# name: test_directory_tree_reloading
+ '''
+
+
+ '''
+# ---
# name: test_disabled_widgets
'''