Merge pull request #4183 from Textualize/batch-async-context-manager

Adds a `Widget.batch` async context manager

Author: rodrigogiraoserrao

CHANGELOG.md

@@ -18,6 +18,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Add attribute `App.animation_level` to control whether animations on that app run or not https://github.com/Textualize/textual/pull/4062
 - Added support for a `TEXTUAL_SCREENSHOT_LOCATION` environment variable to specify the location of an automated screenshot https://github.com/Textualize/textual/pull/4181/
 - Added support for a `TEXTUAL_SCREENSHOT_FILENAME` environment variable to specify the filename of an automated screenshot https://github.com/Textualize/textual/pull/4181/
+- 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

@@ -6,11 +6,13 @@
 
 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,
@@ -3291,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.
 

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)