Anim screenshot fix (#2655)

* fix screenshots in docs

* fix anim

* added wait for animation switch

* remove comment

src/textual/_doc.py

@@ -34,7 +34,12 @@ def format_svg(source, language, css_class, options, md, attrs, **kwargs) -> str
             rows = int(attrs.get("lines", 24))
             columns = int(attrs.get("columns", 80))
             svg = take_svg_screenshot(
-                None, path, press, title, terminal_size=(columns, rows)
+                None,
+                path,
+                press,
+                title,
+                terminal_size=(columns, rows),
+                wait_for_animation=False,
             )
         finally:
             os.chdir(cwd)
@@ -56,6 +61,7 @@ def take_svg_screenshot(
     title: str | None = None,
     terminal_size: tuple[int, int] = (80, 24),
     run_before: Callable[[Pilot], Awaitable[None] | None] | None = None,
+    wait_for_animation: bool = True,
 ) -> str:
     """
 
@@ -68,6 +74,7 @@ def take_svg_screenshot(
         run_before: An arbitrary callable that runs arbitrary code before taking the
             screenshot. Use this to simulate complex user interactions with the app
             that cannot be simulated by key presses.
+        wait_for_animation: Wait for animation to complete before taking screenshot.
 
     Returns:
         An SVG string, showing the content of the terminal window at the time
@@ -109,8 +116,9 @@ def take_svg_screenshot(
             if inspect.isawaitable(result):
                 await result
         await pilot.press(*press)
-        await pilot.wait_for_scheduled_animations()
-        await pilot.pause()
+        if wait_for_animation:
+            await pilot.wait_for_scheduled_animations()
+            await pilot.pause()
         svg = app.export_screenshot(title=title)
 
         app.exit(svg)

src/textual/app.py

@@ -1030,14 +1030,11 @@ class App(Generic[ReturnType], DOMNode):
         app = self
         driver = app._driver
         assert driver is not None
-        await wait_for_idle(0)
         for key in keys:
             if key.startswith("wait:"):
                 _, wait_ms = key.split(":")
                 print(f"(pause {wait_ms}ms)")
                 await asyncio.sleep(float(wait_ms) / 1000)
-                await app._animator.wait_until_complete()
-                await wait_for_idle(0)
             else:
                 if len(key) == 1 and not key.isalnum():
                     key = _character_to_key(key)
@@ -1052,9 +1049,8 @@ class App(Generic[ReturnType], DOMNode):
                 key_event._set_sender(app)
                 driver.send_event(key_event)
                 await wait_for_idle(0)
-
-        await app._animator.wait_until_complete()
-        await wait_for_idle(0)
+                await app._animator.wait_until_complete()
+                await wait_for_idle(0)
 
     @asynccontextmanager
     async def run_test(
@@ -1110,7 +1106,9 @@ class App(Generic[ReturnType], DOMNode):
 
         # Context manager returns pilot object to manipulate the app
         try:
-            yield Pilot(app)
+            pilot = Pilot(app)
+            await pilot._wait_for_screen()
+            yield pilot
         finally:
             # Shutdown the app cleanly
             await app._shutdown()

tests/test_data_table.py

@@ -1009,6 +1009,7 @@ async def test_scrolling_cursor_into_view():
         table.add_column("n")
         table.add_rows([(n,) for n in range(300)])
 
+        await pilot.pause()
         await pilot.press("c")
         await pilot.pause()
         assert table.scroll_y > 100