Merge pull request #2545 from davep/directory-tree-work-in-worker

Load `DirectoryTree` contents in a worker

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ### Changed
 
 - App `title` and `sub_title` attributes can be set to any type https://github.com/Textualize/textual/issues/2521
+- `DirectoryTree` now loads directory contents in a worker https://github.com/Textualize/textual/issues/2456
 - Only a single error will be written by default, unless in dev mode ("debug" in App.features) https://github.com/Textualize/textual/issues/2480
 - Using `Widget.move_child` where the target and the child being moved are the same is now a no-op https://github.com/Textualize/textual/issues/1743
 - Calling `dismiss` on a screen that is not at the top of the stack now raises an exception https://github.com/Textualize/textual/issues/2575

src/textual/widgets/_directory_tree.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
+from asyncio import Queue
 from dataclasses import dataclass
 from pathlib import Path
-from typing import ClassVar, Iterable
+from typing import ClassVar, Iterable, Iterator
 
 from rich.style import Style
 from rich.text import Text, TextType
 
-from ..events import Mount
+from .. import work
 from ..message import Message
 from ..reactive import var
+from ..worker import Worker, WorkerCancelled, WorkerFailed, get_current_worker
 from ._tree import TOGGLE_STYLE, Tree, TreeNode
 
 
@@ -90,7 +92,7 @@ def control(self) -> DirectoryTree:
             """
             return self.tree
 
-    path: var[str | Path] = var["str | Path"](Path("."), init=False)
+    path: var[str | Path] = var["str | Path"](Path("."), init=False, always_update=True)
     """The path that is the root of the directory tree.
 
     Note:
@@ -116,6 +118,7 @@ def __init__(
             classes: A space-separated list of classes, or None for no classes.
             disabled: Whether the directory tree is disabled or not.
         """
+        self._load_queue: Queue[TreeNode[DirEntry]] = Queue()
         super().__init__(
             str(path),
             data=DirEntry(Path(path)),
@@ -126,10 +129,26 @@ def __init__(
         )
         self.path = path
 
+    def _add_to_load_queue(self, node: TreeNode[DirEntry]) -> None:
+        """Add the given node to the load queue.
+
+        Args:
+            node: The node to add to the load queue.
+        """
+        assert node.data is not None
+        node.data.loaded = True
+        self._load_queue.put_nowait(node)
+
     def reload(self) -> None:
         """Reload the `DirectoryTree` contents."""
         self.reset(str(self.path), DirEntry(Path(self.path)))
-        self._load_directory(self.root)
+        # Orphan the old queue...
+        self._load_queue = Queue()
+        # ...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.
+        self._add_to_load_queue(self.root)
 
     def validate_path(self, path: str | Path) -> Path:
         """Ensure that the path is of the `Path` type.
@@ -229,37 +248,115 @@ def filter_paths(self, paths: Iterable[Path]) -> Iterable[Path]:
         """
         return paths
 
-    def _load_directory(self, node: TreeNode[DirEntry]) -> None:
-        """Load the directory contents for a given node.
+    @staticmethod
+    def _safe_is_dir(path: Path) -> bool:
+        """Safely check if a path is a directory.
 
         Args:
-            node: The node to load the directory contents for.
+            path: The path to check.
+
+        Returns:
+            `True` if the path is for a directory, `False` if not.
         """
-        assert node.data is not None
-        node.data.loaded = True
-        directory = sorted(
-            self.filter_paths(node.data.path.iterdir()),
-            key=lambda path: (not path.is_dir(), path.name.lower()),
-        )
-        for path in directory:
+        try:
+            return path.is_dir()
+        except PermissionError:
+            # We may or may not have been looking at a directory, but we
+            # don't have the rights or permissions to even know that. Best
+            # we can do, short of letting the error blow up, is assume it's
+            # not a directory. A possible improvement in here could be to
+            # have a third state which is "unknown", and reflect that in the
+            # tree.
+            return False
+
+    def _populate_node(self, node: TreeNode[DirEntry], content: Iterable[Path]) -> None:
+        """Populate the given tree node with the given directory content.
+
+        Args:
+            node: The Tree node to populate.
+            content: The collection of `Path` objects to populate the node with.
+        """
+        for path in content:
             node.add(
                 path.name,
                 data=DirEntry(path),
-                allow_expand=path.is_dir(),
+                allow_expand=self._safe_is_dir(path),
             )
         node.expand()
 
-    def _on_mount(self, _: Mount) -> None:
-        self._load_directory(self.root)
+    def _directory_content(self, location: Path, worker: Worker) -> Iterator[Path]:
+        """Load the content of a given directory.
+
+        Args:
+            location: The location to load from.
+            worker: The worker that the loading is taking place in.
+
+        Yields:
+            Path: A entry within the location.
+        """
+        try:
+            for entry in location.iterdir():
+                if worker.is_cancelled:
+                    break
+                yield entry
+        except PermissionError:
+            pass
+
+    @work
+    def _load_directory(self, node: TreeNode[DirEntry]) -> list[Path]:
+        """Load the directory contents for a given node.
+
+        Args:
+            node: The node to load the directory contents for.
+
+        Returns:
+            The list of entries within the directory associated with the node.
+        """
+        assert node.data is not None
+        return sorted(
+            self.filter_paths(
+                self._directory_content(node.data.path, get_current_worker())
+            ),
+            key=lambda path: (not self._safe_is_dir(path), path.name.lower()),
+        )
+
+    @work(exclusive=True)
+    async def _loader(self) -> None:
+        """Background loading queue processor."""
+        worker = get_current_worker()
+        while not worker.is_cancelled:
+            # Get the next node that needs loading off the queue. Note that
+            # 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)
+            # Mark this iteration as done.
+            self._load_queue.task_done()
 
     def _on_tree_node_expanded(self, event: Tree.NodeExpanded) -> None:
         event.stop()
         dir_entry = event.node.data
         if dir_entry is None:
             return
-        if dir_entry.path.is_dir():
+        if self._safe_is_dir(dir_entry.path):
             if not dir_entry.loaded:
-                self._load_directory(event.node)
+                self._add_to_load_queue(event.node)
         else:
             self.post_message(self.FileSelected(self, event.node, dir_entry.path))
 
@@ -268,5 +365,5 @@ def _on_tree_node_selected(self, event: Tree.NodeSelected) -> None:
         dir_entry = event.node.data
         if dir_entry is None:
             return
-        if not dir_entry.path.is_dir():
+        if not self._safe_is_dir(dir_entry.path):
             self.post_message(self.FileSelected(self, event.node, dir_entry.path))