improve robustness of AwaitableResponse

nicegui/__init__.py

@@ -1,7 +1,6 @@
 from . import context, elements, run, ui
 from .api_router import APIRouter
 from .app.app import App
-from .awaitable_response import AwaitableResponse
 from .client import Client
 from .nicegui import app
 from .tailwind import Tailwind
@@ -11,7 +10,6 @@ __all__ = [
     'APIRouter',
     'app',
     'App',
-    'AwaitableResponse',
     'Client',
     'context',
     'elements',

nicegui/awaitable_response.py

@@ -1,31 +1,44 @@
 from __future__ import annotations
 
-from asyncio import Task
-from typing import Callable, Optional
+from typing import Callable
 
 from . import background_tasks
 
 
 class AwaitableResponse:
 
-    def __init__(self, fire_and_forget: Optional[Callable], wait_for_result: Optional[Callable]) -> None:
+    def __init__(self, fire_and_forget: Callable, wait_for_result: Callable) -> None:
         """Awaitable Response
 
         This class can be used to run one of two different callables, depending on whether the response is awaited or not.
+        It must be awaited immediately after creation or not at all.
 
         :param fire_and_forget: The callable to run if the response is not awaited.
         :param wait_for_result: The callable to run if the response is awaited.
         """
-        self.fire_and_forget_task: Optional[Task] = \
-            background_tasks.create(self._start(fire_and_forget), name='fire and forget') if fire_and_forget else None
+        self.fire_and_forget = fire_and_forget
         self.wait_for_result = wait_for_result
+        self._is_fired = False
+        self._is_awaited = False
+        background_tasks.create(self._fire(), name='fire')
 
-    async def _start(self, command: Callable) -> None:
-        command()
+    async def _fire(self) -> None:
+        if self._is_awaited:
+            return
+        self._is_fired = True
+        self.fire_and_forget()
 
     def __await__(self):
-        if self.fire_and_forget_task is not None:
-            self.fire_and_forget_task.cancel()
-        if self.wait_for_result is None:
-            raise ValueError('AwaitableResponse has no result to await')
+        if self._is_fired:
+            raise RuntimeError('AwaitableResponse must be awaited immediately after creation or not at all')
+        self._is_awaited = True
         return self.wait_for_result().__await__()
+
+
+class NullResponse(AwaitableResponse):
+
+    def __init__(self) -> None:
+        pass
+
+    def __await__(self):
+        raise RuntimeError('NullResponse cannot be awaited')

nicegui/element.py

@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional,
 from typing_extensions import Self
 
 from . import context, core, events, json, outbox, storage
-from .awaitable_response import AwaitableResponse
+from .awaitable_response import AwaitableResponse, NullResponse
 from .dependencies import Component, Library, register_library, register_vue_component
 from .elements.mixins.visibility import Visibility
 from .event_listener import EventListener
@@ -410,7 +410,7 @@ class Element(Visibility):
         :param args: arguments to pass to the method
         """
         if not core.loop:
-            return AwaitableResponse(None, None)
+            return NullResponse()
         return self.client.run_javascript(f'return runMethod({self.id}, "{name}", {json.dumps(args)})')
 
     def _collect_descendants(self, *, include_self: bool = False) -> List[Element]:

nicegui/elements/scene.py

@@ -4,7 +4,7 @@ from typing import Any, Callable, Dict, List, Optional, Union
 from typing_extensions import Self
 
 from .. import binding
-from ..awaitable_response import AwaitableResponse
+from ..awaitable_response import AwaitableResponse, NullResponse
 from ..dataclasses import KWONLY_SLOTS
 from ..element import Element
 from ..events import (GenericEventArguments, SceneClickEventArguments, SceneClickHit, SceneDragEventArguments,
@@ -124,7 +124,7 @@ class Scene(Element,
         :param args: arguments to pass to the method
         """
         if not self.is_initialized:
-            return AwaitableResponse(None, None)
+            return NullResponse()
         return super().run_method(name, *args)
 
     def _handle_click(self, e: GenericEventArguments) -> None:

tests/test_awaitable_response.py

@@ -0,0 +1,45 @@
+import asyncio
+
+import pytest
+
+from nicegui import core
+from nicegui.awaitable_response import AwaitableResponse, NullResponse
+
+
+async def test_awaitable_response():
+    core.loop = asyncio.get_event_loop()
+    actions = []
+
+    def do_something() -> AwaitableResponse:
+        def fire_and_forget():
+            actions.append('fire_and_forget')
+
+        async def wait_for_result() -> str:
+            actions.append('wait_for_result')
+            return 'result'
+
+        return AwaitableResponse(fire_and_forget, wait_for_result)
+
+    actions.clear()
+    do_something()
+    await asyncio.sleep(0.1)
+    assert actions == ['fire_and_forget']
+
+    actions.clear()
+    result = await do_something()
+    assert result == 'result'
+    assert actions == ['wait_for_result']
+
+    actions.clear()
+    result = do_something()
+    await asyncio.sleep(0.1)
+    with pytest.raises(RuntimeError, match='AwaitableResponse must be awaited immediately after creation or not at all'):
+        await result
+    assert actions == ['fire_and_forget']
+
+
+async def test_null_response():
+    NullResponse()
+
+    with pytest.raises(RuntimeError, match='NullResponse cannot be awaited'):
+        await NullResponse()