Adds a Widget.batch async context manager

CHANGELOG.md

@@ -11,6 +11,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 - Added support for environment variable `TEXTUAL_ANIMATIONS` to control what animations Textual displays https://github.com/Textualize/textual/pull/4062
 - Add attribute `App.animation_level` to control whether animations on that app run or not https://github.com/Textualize/textual/pull/4062
+- Added an `asyncio` lock attribute `Widget.lock` to be used to synchronize widget state https://github.com/Textualize/textual/issues/4134
+- `Widget.remove_children` now accepts a CSS selector to specify which children to remove https://github.com/Textualize/textual/pull/4183
+- `Widget.batch` combines widget locking and app update batching https://github.com/Textualize/textual/pull/4183
 
 ## [0.51.0] - 2024-02-15
 

src/textual/widget.py

@@ -4,13 +4,15 @@
 
 from __future__ import annotations
 
-from asyncio import create_task, wait
+from asyncio import Lock, create_task, wait
 from collections import Counter
+from contextlib import asynccontextmanager
 from fractions import Fraction
 from itertools import islice
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
+    AsyncGenerator,
     Awaitable,
     ClassVar,
     Collection,
@@ -53,6 +55,8 @@
 from .await_remove import AwaitRemove
 from .box_model import BoxModel
 from .cache import FIFOCache
+from .css.match import match
+from .css.parse import parse_selectors
 from .css.query import NoMatches, WrongType
 from .css.scalar import ScalarOffset
 from .dom import DOMNode, NoScreen
@@ -78,6 +82,7 @@
 
 if TYPE_CHECKING:
     from .app import App, ComposeResult
+    from .css.query import QueryType
     from .message_pump import MessagePump
     from .scrollbar import (
         ScrollBar,
@@ -373,6 +378,14 @@ def __init__(
         if self.BORDER_SUBTITLE:
             self.border_subtitle = self.BORDER_SUBTITLE
 
+        self.lock = Lock()
+        """`asyncio` lock to be used to synchronize the state of the widget.
+
+        Two different tasks might call methods on a widget at the same time, which
+        might result in a race condition.
+        This can be fixed by adding `async with widget.lock:` around the method calls.
+        """
+
     virtual_size: Reactive[Size] = Reactive(Size(0, 0), layout=True)
     """The virtual (scrollable) [size][textual.geometry.Size] of the widget."""
 
@@ -3283,15 +3296,42 @@ def remove(self) -> AwaitRemove:
         await_remove = self.app._remove_nodes([self], self.parent)
         return await_remove
 
-    def remove_children(self) -> AwaitRemove:
-        """Remove all children of this Widget from the DOM.
+    def remove_children(self, selector: str | type[QueryType] = "*") -> AwaitRemove:
+        """Remove the immediate children of this Widget from the DOM.
+
+        Args:
+            selector: A CSS selector to specify which direct children to remove.
 
         Returns:
-            An awaitable object that waits for the children to be removed.
+            An awaitable object that waits for the direct children to be removed.
         """
-        await_remove = self.app._remove_nodes(list(self.children), self)
+        if not isinstance(selector, str):
+            selector = selector.__name__
+        parsed_selectors = parse_selectors(selector)
+        children_to_remove = [
+            child for child in self.children if match(parsed_selectors, child)
+        ]
+        await_remove = self.app._remove_nodes(children_to_remove, self)
         return await_remove
 
+    @asynccontextmanager
+    async def batch(self) -> AsyncGenerator[None, None]:
+        """Async context manager that combines widget locking and update batching.
+
+        Use this async context manager whenever you want to acquire the widget lock and
+        batch app updates at the same time.
+
+        Example:
+            ```py
+            async with container.batch():
+                await container.remove_children(Button)
+                await container.mount(Label("All buttons are gone."))
+            ```
+        """
+        async with self.lock:
+            with self.app.batch_update():
+                yield
+
     def render(self) -> RenderableType:
         """Get text or Rich renderable for this widget.
 

src/textual/widgets/_tree.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from asyncio import Lock
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar, Generic, Iterable, NewType, TypeVar, cast
 
@@ -618,8 +617,6 @@ def __init__(
         self._line_cache: LRUCache[LineCacheKey, Strip] = LRUCache(1024)
         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)
 

tests/test_widget_removing.py

@@ -141,6 +141,63 @@ async def test_widget_remove_children_container():
         assert len(container.children) == 0
 
 
+async def test_widget_remove_children_with_star_selector():
+    app = ExampleApp()
+    async with app.run_test():
+        container = app.query_one(Vertical)
+
+        # 6 labels in total, with 5 of them inside the container.
+        assert len(app.query(Label)) == 6
+        assert len(container.children) == 5
+
+        await container.remove_children("*")
+
+        # The labels inside the container are gone, and the 1 outside remains.
+        assert len(app.query(Label)) == 1
+        assert len(container.children) == 0
+
+
+async def test_widget_remove_children_with_string_selector():
+    app = ExampleApp()
+    async with app.run_test():
+        container = app.query_one(Vertical)
+
+        # 6 labels in total, with 5 of them inside the container.
+        assert len(app.query(Label)) == 6
+        assert len(container.children) == 5
+
+        await app.screen.remove_children("Label")
+
+        # Only the Screen > Label widget is gone, everything else remains.
+        assert len(app.query(Button)) == 1
+        assert len(app.query(Vertical)) == 1
+        assert len(app.query(Label)) == 5
+
+
+async def test_widget_remove_children_with_type_selector():
+    app = ExampleApp()
+    async with app.run_test():
+        assert len(app.query(Button)) == 1  # Sanity check.
+        await app.screen.remove_children(Button)
+        assert len(app.query(Button)) == 0
+
+
+async def test_widget_remove_children_with_selector_does_not_leak():
+    app = ExampleApp()
+    async with app.run_test():
+        container = app.query_one(Vertical)
+
+        # 6 labels in total, with 5 of them inside the container.
+        assert len(app.query(Label)) == 6
+        assert len(container.children) == 5
+
+        await container.remove_children("Label")
+
+        # The labels inside the container are gone, and the 1 outside remains.
+        assert len(app.query(Label)) == 1
+        assert len(container.children) == 0
+
+
 async def test_widget_remove_children_no_children():
     app = ExampleApp()
     async with app.run_test():
@@ -154,3 +211,17 @@ async def test_widget_remove_children_no_children():
         assert (
             count_before == count_after
         )  # No widgets have been removed, since Button has no children.
+
+
+async def test_widget_remove_children_no_children_match_selector():
+    app = ExampleApp()
+    async with app.run_test():
+        container = app.query_one(Vertical)
+        assert len(container.query("Button")) == 0  # Sanity check.
+
+        count_before = len(app.query("*"))
+        container_children_before = list(container.children)
+        await container.remove_children("Button")
+
+        assert count_before == len(app.query("*"))
+        assert container_children_before == list(container.children)