Make stream() async so is non-blocking

Author: cpsievert

shiny/api-examples/MarkdownStream/app-core.py

@@ -1,8 +1,8 @@
-import time
+import asyncio
 
 import requests
 
-from shiny import App, ui
+from shiny import App, ui, reactive
 
 app_ui = ui.page_fluid(
     ui.card(
@@ -23,13 +23,16 @@ def server(input, output, session):
     readme_chunks = readme.text.replace("\n", " \n ").split(" ")
 
     # Generate words from the README.md file (with a small delay)
-    def chunk_generator():
+    async def chunk_generator():
         for chunk in readme_chunks:
-            time.sleep(0.02)
+            await asyncio.sleep(0.02)
             yield chunk + " "
 
     md = ui.MarkdownStream("shiny-readme")
-    md.stream(chunk_generator())
+
+    @reactive.effect
+    async def _():
+        await md.stream(chunk_generator())
 
 
 app = App(app_ui, server)

shiny/api-examples/MarkdownStream/app-express.py

@@ -1,7 +1,8 @@
-import time
+import asyncio
 
 import requests
 
+from shiny import reactive
 from shiny.express import session, ui
 
 ui.page_opts(full_width=True)
@@ -14,10 +15,10 @@
 
 
 # Generate words from the README.md file (with a small delay)
-def chunk_generator():
+async def chunk_generator():
     for chunk in readme_chunks:
         if not session.is_stub_session():
-            time.sleep(0.02)
+            await asyncio.sleep(0.02)
         yield chunk + " "
 
 
@@ -27,4 +28,9 @@ def chunk_generator():
     ui.card_header("Shiny README.md")
     md.ui()
 
-md.stream(chunk_generator())
+md.ui(height=300)
+
+
+@reactive.effect
+async def _():
+    await md.stream(chunk_generator())

shiny/ui/_markdown_stream.py

@@ -1,9 +1,9 @@
-from contextlib import contextmanager
-from typing import Iterable, Literal, Union
+from contextlib import asynccontextmanager
+from typing import AsyncIterable, Iterable, Literal, Union
 
 from htmltools import css
 
-from .. import reactive
+from .. import _utils, reactive
 from .._docstring import add_example
 from .._typing_extensions import TypedDict
 from ..session import require_active_session
@@ -125,7 +125,11 @@ def ui(
             height=height,
         )
 
-    def stream(self, content: Iterable[str], clear: bool = True):
+    async def stream(
+        self,
+        content: Iterable[str] | AsyncIterable[str],
+        clear: bool = True,
+    ):
         """
         Stream content into the markdown stream.
 
@@ -138,13 +142,15 @@ def stream(self, content: Iterable[str], clear: bool = True):
             Whether to clear the existing content before streaming the new content.
         """
 
+        content = _utils.wrap_async_iterable(content)
+
         @reactive.extended_task
         async def _task():
             if clear:
-                self._replace("")
-            with self._streaming_dot():
-                for c in content:
-                    self._append(c)
+                await self._replace("")
+            async with self._streaming_dot():
+                async for c in content:
+                    await self._append(c)
 
         _task()
 
@@ -157,31 +163,31 @@ async def _handle_error():
                 await self._raise_exception(e)
             _handle_error.destroy()  # type: ignore
 
-    def _append(self, content: str):
+    async def _append(self, content: str):
         msg: ContentMessage = {
             "id": self.id,
             "content": content,
             "operation": "append",
         }
 
-        self._send_custom_message(msg)
+        await self._send_custom_message(msg)
 
-    def _replace(self, content: str):
+    async def _replace(self, content: str):
         msg: ContentMessage = {
             "id": self.id,
             "content": content,
             "operation": "replace",
         }
 
-        self._send_custom_message(msg)
+        await self._send_custom_message(msg)
 
-    @contextmanager
-    def _streaming_dot(self):
+    @asynccontextmanager
+    async def _streaming_dot(self):
         start: isStreamingMessage = {
             "id": self.id,
             "isStreaming": True,
         }
-        self._send_custom_message(start)
+        await self._send_custom_message(start)
 
         try:
             yield
@@ -190,7 +196,7 @@ def _streaming_dot(self):
                 "id": self.id,
                 "isStreaming": False,
             }
-            self._send_custom_message(end)
+            await self._send_custom_message(end)
 
     async def _raise_exception(self, e: BaseException):
         if self.on_error == "unhandled":
@@ -200,12 +206,12 @@ async def _raise_exception(self, e: BaseException):
             msg = f"Error in MarkdownStream('{self.id}'): {str(e)}"
             raise NotifyException(msg, sanitize=sanitize) from e
 
-    def _send_custom_message(self, msg: Union[ContentMessage, isStreamingMessage]):
+    async def _send_custom_message(
+        self, msg: Union[ContentMessage, isStreamingMessage]
+    ):
         if self._session.is_stub_session():
             return
-        self._session._send_message_sync(
-            {"custom": {"shinyMarkdownStreamMessage": msg}}
-        )
+        await self._session.send_custom_message("shinyMarkdownStreamMessage", msg)
 
 
 @add_example()