Fix on_load event to support async functions (#1300) (#1323)

Co-authored-by: Claude <[email protected]>

Author: richard-to

mesop/examples/testing/__init__.py

@@ -1,3 +1,6 @@
+from mesop.examples.testing import (
+  async_onload as async_onload,
+)
 from mesop.examples.testing import (
   click_is_target as click_is_target,
 )

mesop/examples/testing/async_onload.py

@@ -0,0 +1,45 @@
+"""
+Test script to verify async on_load support.
+This reproduces the issue from GitHub issue #1300.
+"""
+
+import asyncio
+
+import mesop as me
+
+
+@me.stateclass
+class State:
+  message: str = ""
+
+
+async def async_on_load(e: me.LoadEvent):
+  """Async generator on_load handler."""
+  state = me.state(State)
+  state.message = "Loading..."
+  yield
+
+  # Simulate async operation
+  await asyncio.sleep(0.1)
+
+  state.message = "Loaded from async generator!"
+  yield
+
+
+async def async_coroutine_on_load(e: me.LoadEvent):
+  """Async coroutine on_load handler (no yield)."""
+  state = me.state(State)
+  await asyncio.sleep(0.1)
+  state.message = "Loaded from coroutine!"
+
+
+@me.page(path="/async_gen", on_load=async_on_load)
+def page_async_gen():
+  state = me.state(State)
+  me.text(f"Async Generator: {state.message}")
+
+
+@me.page(path="/async_coro", on_load=async_coroutine_on_load)
+def page_async_coro():
+  state = me.state(State)
+  me.text(f"Async Coroutine: {state.message}")

mesop/runtime/context.py

@@ -1,4 +1,3 @@
-import asyncio
 import copy
 import threading
 import types
@@ -18,6 +17,7 @@
   MesopException,
 )
 from mesop.server.state_session import state_session
+from mesop.utils.async_utils import run_async_generator, run_coroutine
 
 T = TypeVar("T")
 
@@ -363,9 +363,9 @@ def run_event_handler(
       result = handler(payload)
       if result is not None:
         if isinstance(result, types.AsyncGeneratorType):
-          yield from _run_async_generator(result)
+          yield from run_async_generator(result)
         elif isinstance(result, types.CoroutineType):
-          yield _run_coroutine(result)
+          yield run_coroutine(result)
         else:
           yield from result
       else:
@@ -374,26 +374,3 @@ def run_event_handler(
       raise MesopException(
         f"Unknown handler id: {event.handler_id} from event {event}"
       )
-
-
-def _run_async_generator(agen: types.AsyncGeneratorType[None, None]):
-  loop = _get_or_create_event_loop()
-  try:
-    while True:
-      yield loop.run_until_complete(agen.__anext__())
-  except StopAsyncIteration:
-    pass
-
-
-def _run_coroutine(coroutine: types.CoroutineType):
-  loop = _get_or_create_event_loop()
-  return loop.run_until_complete(coroutine)
-
-
-def _get_or_create_event_loop():
-  try:
-    return asyncio.get_running_loop()
-  except RuntimeError:
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    return loop

mesop/runtime/runtime.py

@@ -1,3 +1,4 @@
+from collections.abc import AsyncGenerator, Coroutine
 from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, Callable, Generator, Type, TypeVar, cast
@@ -23,7 +24,13 @@ class EmptyState:
   pass
 
 
-OnLoadHandler = Callable[[LoadEvent], None | Generator[None, None, None]]
+OnLoadHandler = Callable[
+  [LoadEvent],
+  None
+  | Generator[None, None, None]
+  | AsyncGenerator[None, None]
+  | Coroutine[Any, Any, None],
+]
 
 
 @dataclass(kw_only=True)

mesop/server/server.py

@@ -2,6 +2,7 @@
 import logging
 import secrets
 import threading
+import types
 from typing import Generator, Sequence
 
 from flask import (
@@ -36,6 +37,7 @@
   prefix_base_url,
   serialize,
 )
+from mesop.utils.async_utils import run_async_generator, run_coroutine
 from mesop.utils.url_utils import remove_url_query_param
 from mesop.warn import warn
 
@@ -44,6 +46,18 @@
 logger = logging.getLogger(__name__)
 
 
+def _process_on_load_result(result) -> Generator[None, None, None]:
+  """Process on_load result, handling sync generators, async generators, and coroutines."""
+  if result is not None:
+    if isinstance(result, types.AsyncGeneratorType):
+      yield from run_async_generator(result)
+    elif isinstance(result, types.CoroutineType):
+      yield run_coroutine(result)
+    else:
+      # Regular generator
+      yield from result
+
+
 def configure_flask_app(
   *, prod_mode: bool = True, exceptions_to_propagate: Sequence[type] = ()
 ) -> Flask:
@@ -172,9 +186,9 @@ def generate_data(ui_request: pb.UiRequest) -> Generator[str, None, None]:
             )
           )
           # on_load is a generator function then we need to iterate through
-          # the generator object.
+          # the generator object. This also handles async generators and coroutines.
           if result:
-            for _ in result:
+            for _ in _process_on_load_result(result):
               yield from render_loop(path=ui_request.path, init_request=True)
               runtime().context().set_previous_node_from_current_node()
               runtime().context().reset_current_node()
@@ -277,9 +291,9 @@ def run_page_load(*, path: str):
     assert page_config and page_config.on_load
     result = page_config.on_load(LoadEvent(path=path))
     # on_load is a generator function then we need to iterate through
-    # the generator object.
+    # the generator object. This also handles async generators and coroutines.
     if result:
-      for _ in result:
+      for _ in _process_on_load_result(result):
         yield from render_loop(path=path, init_request=True)
         runtime().context().set_previous_node_from_current_node()
         runtime().context().reset_current_node()

mesop/tests/e2e/async_onload_test.ts

@@ -0,0 +1,19 @@
+import {test, expect} from '@playwright/test';
+
+test.describe('async onload events', () => {
+  test('async generator', async ({page}) => {
+    await page.goto('/async_gen');
+
+    await expect(
+      page.getByText('Async Generator: Loaded from async generator!'),
+    ).toBeVisible();
+  });
+
+  test('async coroutine', async ({page}) => {
+    await page.goto('/async_coro');
+
+    await expect(
+      page.getByText('Async Coroutine: Loaded from coroutine!'),
+    ).toBeVisible();
+  });
+});

mesop/utils/async_utils.py

@@ -0,0 +1,51 @@
+"""Utilities for handling async generators and coroutines in Mesop event handlers."""
+
+import asyncio
+import types
+from typing import Generator
+
+
+def run_async_generator(
+  agen: types.AsyncGeneratorType[None, None],
+) -> Generator[None, None, None]:
+  """Run an async generator by iterating through it using an event loop.
+
+  Args:
+    agen: The async generator to run
+
+  Yields:
+    None for each iteration of the async generator
+  """
+  loop = get_or_create_event_loop()
+  try:
+    while True:
+      yield loop.run_until_complete(agen.__anext__())
+  except StopAsyncIteration:
+    pass
+
+
+def run_coroutine(coroutine: types.CoroutineType):
+  """Run a coroutine using an event loop and return its result.
+
+  Args:
+    coroutine: The coroutine to run
+
+  Returns:
+    The result of the coroutine
+  """
+  loop = get_or_create_event_loop()
+  return loop.run_until_complete(coroutine)
+
+
+def get_or_create_event_loop():
+  """Get the current event loop or create a new one if none exists.
+
+  Returns:
+    The asyncio event loop
+  """
+  try:
+    return asyncio.get_running_loop()
+  except RuntimeError:
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    return loop