Repository: Textualize/textual
Pull request: Adds a `Widget.batch` async context manager
The `Widget.batch` context manager locks the widget and batches app updates too.
This fixes #4133.

This PR depends on #4139 and shouldn't be merged while #4139 doesn't get the OK to be merged.

Problem statement:
Add `remove` attribute to `mount` and `mount_all`
It's common to remove some widgets and add some other widgets. To make this easier, I think we should add a `remove` attribute to the mount methods which accepts a selector and removes those widgets prior to mounting.

The remove + mount should be atomic, to avoid flicker (might need to be wrapped in `batch_update`).

API would look something like this:

```python
# Replace all MyWidgets with a single MyWidget
await self.mount(MyWidget("Foo"), remove="MyWidget")

# Remove all children and replace with a label
await self.mount(Label("Goodbye"), remove="*")
```

Hints from issue discussion:
I think of the method `mount` as a method that “builds”, “adds”, “constructs”.
Adding a kwd argument that goes in the opposite direction feels a bit clunky.

Wouldn't it make more sense to grow a method `replace(selector, new_widget)`?
The thing is that it is not replacing. It's not like we can guarantee the widgets will be mounted in the same place as the "remove" selector.
In that case, call it `remove_then_mount`.
Or, better yet, let `Widget.remove_children` grow the `selector` parameter that is `"*"` by default but that you can change.
I think

```py
self.remove_children(selector="...")
self.mount(widget)
```

looks much better than

```py
self.mount(widget, remove="...")
```

The thing with `remove` inside `mount` is that people will think that the parameter `remove` is somehow _related_ to the mounting when it's completely orthogonal.
Two methods aren't atomic. Which means you would have to wrap both calls in an asyncio Lock, and await them both. You would also need to use `batch_update` to avoid flicker. So now your code will look like this:

```python
async with self.lock:
    with self.app.batch_update():
        await self.remove_children("*")
        await self.mount(widget)
```

That's too much to ask of devs.

The `remove` attribute on `mount` is still my favorite. People misunderstanding the extra parameter can be alleviated by docs.
If you prefer the `mount(..., remove=selector)` over `remove_then_mount`, sure.
After a discussion in the office. We could combine the locking and batch update in to a single context manager.

```python
async with self.batch():
    await self.remove_children("*")
    await self.mount(widget)
```

That feels a little less clumsy, but it has its own flaw: Namely, that it requires async. You wouldn't be able to have optional awaitables.

Commit message: Merge branch 'main' into batch-async-context-manager

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 204fcb5f42..1aac44edff 100644
--- a/CHANGELOG.md
+++ b/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
 
diff --git a/src/textual/widget.py b/src/textual/widget.py
index 4a3f6c12fe..aa56faae2f 100644
--- a/src/textual/widget.py
+++ b/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,
@@ -78,6 +80,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 +376,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 +3294,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] | None = None
+    ) -> AwaitRemove:
+        """Remove the children of this Widget from the DOM.
+
+        Args:
+            selector: A CSS selector to specify which children to remove.
 
         Returns:
             An awaitable object that waits for the children to be removed.
         """
-        await_remove = self.app._remove_nodes(list(self.children), self)
+        if isinstance(selector, str) or selector is None:
+            children_to_remove = self.query(selector)
+        else:
+            children_to_remove = self.query(selector.__name__)
+        await_remove = self.app._remove_nodes(list(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.
 
diff --git a/src/textual/widgets/_tree.py b/src/textual/widgets/_tree.py
index 3e0a2d8cb4..f218ae86f2 100644
--- a/src/textual/widgets/_tree.py
+++ b/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)
 
diff --git a/tests/test_widget_removing.py b/tests/test_widget_removing.py
index ddb9dae16a..d81867e5c7 100644
--- a/tests/test_widget_removing.py
+++ b/tests/test_widget_removing.py
@@ -141,6 +141,79 @@ 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")
+
+        # All labels are gone but the button and the container remain.
+        assert len(app.query(Button)) == 1
+        assert len(app.query(Vertical)) == 1
+        assert len(app.query(Label)) == 0
+
+
+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_with_compound_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("Vertical > 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 +227,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)


Code review comments:
diff_hunk: @@ -3283,15 +3294,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(
path: src/textual/widget.py
text: @user1:
This method looks like it applies to children and their descendants? Basically equivalent to `self.query(selector).remove()`, so there doesn't seem much point in this existing.

What we could probably use is a method to remove only direct children, which would better serve the use case of removing and mounting in a batch.



@author:
Understood. See 808f536.