Yield to the event loop while building Markdown and composing widgets

src/textual/app.py

@@ -94,7 +94,7 @@
 from textual.await_remove import AwaitRemove
 from textual.binding import Binding, BindingsMap, BindingType, Keymap
 from textual.command import CommandListItem, CommandPalette, Provider, SimpleProvider
-from textual.compose import compose
+from textual.compose import compose_async
 from textual.content import Content
 from textual.css.errors import StylesheetError
 from textual.css.query import NoMatches
@@ -3567,7 +3567,7 @@ async def _on_compose(self) -> None:
         _rich_traceback_omit = True
         self._compose_screen = self.screen
         try:
-            widgets = [*self.screen._nodes, *compose(self)]
+            widgets = [*self.screen._nodes, *await compose_async(self)]
         except TypeError as error:
             raise TypeError(
                 f"{self!r} compose() method returned an invalid result; {error}"
@@ -3591,7 +3591,7 @@ async def recompose(self) -> None:
         try:
             async with self.screen.batch():
                 await self.screen.query("*").exclude(".-textual-system").remove()
-                await self.screen.mount_all(compose(self))
+                await self.screen.mount_all(await compose_async(self))
         except ScreenStackError:
             pass
 

src/textual/compose.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import asyncio
+from types import GeneratorType
+from typing import TYPE_CHECKING, Generator
 
 if TYPE_CHECKING:
     from textual.app import App, ComposeResult
@@ -35,6 +37,47 @@ def add_key(key:str) -> ComposeResult:
     Returns:
         A list of widgets.
     """
+    composition = _compose(node, compose_result)
+    try:
+        while True:
+            next(composition)
+    except StopIteration as completed:
+        return completed.value
+    finally:
+        composition.close()
+
+
+async def compose_async(
+    node: App | Widget, compose_result: ComposeResult | None = None
+) -> list[Widget]:
+    """Collect the same declaration without occupying the loop between children.
+
+    Children stay private until the complete result enters registration. A
+    suspended ``with`` belongs to this generator, not another pump's compose;
+    lend its original stacks only while that generator is executing.
+    """
+    app = node.app
+    composition = _compose(node, compose_result)
+    try:
+        while True:
+            next(composition)
+            compose_stack = app._compose_stacks.pop()
+            composed = app._composed.pop()
+            try:
+                await asyncio.sleep(0)
+            finally:
+                app._compose_stacks.append(compose_stack)
+                app._composed.append(composed)
+    except StopIteration as completed:
+        return completed.value
+    finally:
+        composition.close()
+
+
+def _compose(
+    node: App | Widget, compose_result: ComposeResult | None
+) -> Generator[None, None, list[Widget]]:
+    """Own declaration traversal and validation for both acquisition paths."""
     _rich_traceback_omit = True
     from textual.widget import MountError, Widget
 
@@ -44,11 +87,12 @@ def add_key(key:str) -> ComposeResult:
     composed: list[Widget] = []
     app._compose_stacks.append(compose_stack)
     app._composed.append(composed)
-    iter_compose = iter(
-        compose_result if compose_result is not None else node.compose()
-    )
-    is_generator = hasattr(iter_compose, "throw")
+    iter_compose = None
     try:
+        iter_compose = iter(
+            compose_result if compose_result is not None else node.compose()
+        )
+        is_generator = hasattr(iter_compose, "throw")
         while True:
             try:
                 child = next(iter_compose)
@@ -90,10 +134,17 @@ def add_key(key:str) -> ComposeResult:
                         raise
             else:
                 nodes.append(child)
+            yield
         if composed:
             nodes.extend(composed)
             composed.clear()
     finally:
-        app._compose_stacks.pop()
-        app._composed.pop()
+        try:
+            # Cancellation may occur while a declaration's with-block is open.
+            # Close it with its own stacks installed before releasing them.
+            if isinstance(iter_compose, GeneratorType):
+                iter_compose.close()
+        finally:
+            app._compose_stacks.pop()
+            app._composed.pop()
     return nodes

src/textual/widget.py

@@ -70,7 +70,7 @@
 from textual.box_model import BoxModel
 from textual.cache import FIFOCache, LRUCache
 from textual.color import Color
-from textual.compose import compose
+from textual.compose import compose, compose_async
 from textual.content import Content, ContentType
 from textual.css.match import match
 from textual.css.parse import parse_selectors
@@ -1843,7 +1843,7 @@ async def recompose(self) -> None:
         async with self.batch():
             await self.query_children("*").exclude(".-textual-system").remove()
             if self.is_attached:
-                compose_nodes = compose(self)
+                compose_nodes = await compose_async(self)
                 await self.mount_all(compose_nodes)
 
     def _post_register(self, app: App) -> None:
@@ -5072,7 +5072,7 @@ async def _on_compose(self, event: events.Compose) -> None:
 
     async def _compose(self) -> None:
         try:
-            widgets = [*self._pending_children, *compose(self)]
+            widgets = [*self._pending_children, *await compose_async(self)]
             self._pending_children.clear()
         except TypeError as error:
             raise TypeError(

src/textual/widgets/_markdown.py

@@ -8,7 +8,7 @@
 from functools import partial
 from pathlib import Path, PurePath
 from types import MethodType
-from typing import Callable, Iterable, Optional
+from typing import AsyncIterator, Callable, Iterable, Optional
 from urllib.parse import unquote
 
 from markdown_it import MarkdownIt
@@ -732,6 +732,7 @@ async def _update_rows(self, updated_rows: list[list[Content]]) -> None:
                         classes=f"row{row_index} cell",
                     ).with_tooltip(cell)
                 )
+                await asyncio.sleep(0)
         self.last_row = row_index
         await self.mount_all(new_cells)
 
@@ -1329,7 +1330,7 @@ def unhandled_token(self, token: Token) -> MarkdownBlock | None:
         """
         return None
 
-    def _parse_markdown(self, tokens: Iterable[Token]) -> Iterable[MarkdownBlock]:
+    async def _parse_markdown(self, tokens: Iterable[Token]) -> AsyncIterator[MarkdownBlock]:
         """Create a stream of MarkdownBlock widgets from markdown.
 
         Args:
@@ -1345,6 +1346,9 @@ def _parse_markdown(self, tokens: Iterable[Token]) -> Iterable[MarkdownBlock]:
         get_block_class = self.get_block_class
 
         for token in tokens:
+            # One root may own a whole nested list or table. Resume other pumps
+            # between token-owned construction, not after a fixed root count.
+            await asyncio.sleep(0)
             token_type = token.type
             if token_type == "heading_open":
                 stack_append(get_block_class(token.tag)(self, token))
@@ -1418,23 +1422,6 @@ def _parse_markdown(self, tokens: Iterable[Token]) -> Iterable[MarkdownBlock]:
                     else:
                         yield external
 
-    def _build_from_source(self, markdown: str) -> list[MarkdownBlock]:
-        """Build blocks from markdown source.
-
-        Args:
-            markdown: A Markdown document, or partial document.
-
-        Returns:
-            A list of MarkdownBlock instances.
-        """
-        parser = (
-            MarkdownIt("gfm-like")
-            if self._parser_factory is None
-            else self._parser_factory()
-        )
-        tokens = parser.parse(markdown)
-        return list(self._parse_markdown(tokens))
-
     async def _parse_tokens(
         self, parser: MarkdownIt, markdown: str, *, use_thread: bool
     ) -> list[Token] | None:
@@ -1470,41 +1457,25 @@ def update(self, markdown: str) -> AwaitComplete:
         self._table_of_contents = None
 
         async def await_update() -> None:
-            """Update in batches."""
-            BATCH_SIZE = 200
-            batch: list[MarkdownBlock] = []
+            """Construct and complete each original root before starting the next."""
 
             # Lock so that you can't update with more than one document simultaneously
             async with self.lock:
                 tokens = await self._parse_tokens(parser, markdown, use_thread=True)
                 if tokens is None:
                     return
 
-                # Remove existing blocks for the first batch only
+                # Replace existing blocks with the first completed root only.
                 removed: bool = False
 
-                async def mount_batch(batch: list[MarkdownBlock]) -> None:
-                    """Mount a single match of blocks.
-
-                    Args:
-                        batch: A list of blocks to mount.
-                    """
-                    nonlocal removed
+                async for block in self._parse_markdown(tokens):
                     if removed:
-                        await self.mount_all(batch)
+                        await self.mount(block)
                     else:
                         async with self.batch():
                             await markdown_block.remove()
-                            await self.mount_all(batch)
+                            await self.mount(block)
                         removed = True
-
-                for block in self._parse_markdown(tokens):
-                    batch.append(block)
-                    if len(batch) == BATCH_SIZE:
-                        await mount_batch(batch)
-                        batch.clear()
-                if batch:
-                    await mount_batch(batch)
                 if not removed:
                     await markdown_block.remove()
 
@@ -1559,7 +1530,7 @@ async def await_append() -> None:
                         self._last_parsed_line += token.map[0]
                         break
 
-                new_blocks = list(self._parse_markdown(tokens))
+                new_blocks = [block async for block in self._parse_markdown(tokens)]
                 any_headers = any(
                     isinstance(block, MarkdownHeader) for block in new_blocks
                 )
@@ -1581,8 +1552,8 @@ async def await_append() -> None:
                         else:
                             new_blocks = new_blocks[1:]
 
-                    if new_blocks:
-                        await self.mount_all(new_blocks)
+                    for block in new_blocks:
+                        await self.mount(block)
 
                 if any_headers:
                     self._table_of_contents = None

tests/test_markdown.py

@@ -43,6 +43,71 @@ def compose(self) -> ComposeResult:
         yield FussyMarkdown(self._markdown)
 
 
+@pytest.mark.parametrize("kind", ["list", "table"])
+async def test_driver_wheel_completes_during_nested_markdown_construction(kind):
+    """A single root must not retain the UI turn for its whole token tree."""
+    from textual import events
+    from textual.containers import VerticalScroll
+    from textual.widgets import Static
+
+    delivered = asyncio.Event()
+    construction = []
+    at_delivery = []
+    rows = 40
+    source = (
+        "\n".join(f"- row {index}\n  - nested {index}" for index in range(rows))
+        if kind == "list" else
+        "| Heading |\n| --- |\n" + "\n".join(f"| row {index} |" for index in range(rows))
+    )
+
+    class Document(Markdown):
+        def get_block_class(self, name):
+            if name in ("paragraph_open", "td_open"):
+                construction.append(name)
+                if len(construction) == 1:
+                    reader = self.app.query_one("#reader", VerticalScroll)
+                    point = reader.scrollable_content_region.offset
+                    target, _ = self.app.get_widget_at(point.x + 1, point.y + 1)
+                    assert target is reader or reader in target.ancestors
+                    self.app._driver.process_message(events.MouseScrollDown(
+                        None, point.x + 1, point.y + 1, 0, 0, 0, False, False, False
+                    ))
+            return super().get_block_class(name)
+
+    class ReadingApp(App):
+        def compose(self):
+            with VerticalScroll(id="reader"):
+                yield Static("Existing body\n" * 60, id="reader-body")
+                yield Document(id="document")
+
+        async def on_event(self, event):
+            ingress = isinstance(event, events.MouseScrollDown) and not event.is_forwarded
+            result = await super().on_event(event)
+            if ingress:
+                at_delivery.append((len(construction), event.widget))
+                delivered.set()
+            return result
+
+    app = ReadingApp()
+    async with app.run_test(size=(60, 16)) as pilot:
+        await pilot.pause()
+        assert app.query_one("#reader", VerticalScroll).allow_vertical_scroll
+        document = app.query_one(Document)
+        publication = document.update(source)
+        try:
+            await asyncio.wait_for(delivered.wait(), 3)
+            assert not publication.is_done
+            reader = app.query_one("#reader", VerticalScroll)
+            assert reader.scroll_y > 0, (at_delivery, reader.max_scroll_y, reader.region)
+        finally:
+            await publication
+        assert at_delivery[0][0] < len(construction)
+        assert len(document.children) == 1
+        assert all(child.is_mounted for child in document.walk_children())
+        assert document.source == source
+        assert app._compose_stacks == [] and app._composed == []
+
+
 @pytest.mark.parametrize("parked", [False, True])
 async def test_stream_empty_stop_joins_without_publication(parked):
     appends = []