moved update-after-await code

nicegui/auto_context.py

@@ -1,6 +1,5 @@
 import asyncio
-from contextlib import contextmanager
-from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Generator, List, Optional, Union
+from typing import Any, Coroutine, Generator, List
 
 from . import globals
 from .task_logger import create_task
@@ -41,3 +40,28 @@ class within_view(object):
         if len(self.view) != self.child_count:
             self.child_count = len(self.view)
             create_task(self.view.update())
+
+
+class update_after_await:
+
+    def __init__(self, coro: Coroutine, context):
+        self.coro = coro
+        self.context = context
+        self.context.lazy_update()
+
+    def __await__(self) -> Generator[Any, None, Any | None]:
+        coro_iter = self.coro.__await__()
+        iter_send, iter_throw = coro_iter.send, coro_iter.throw
+        send, message = iter_send, None
+        while True:
+            try:
+                signal = send(message)
+                self.context.lazy_update()
+            except StopIteration as err:
+                return err.value
+            else:
+                send = iter_send
+            try:
+                message = yield signal
+            except BaseException as err:
+                send, message = iter_throw, err

nicegui/page.py

@@ -6,7 +6,7 @@ import time
 import types
 import uuid
 from functools import wraps
-from typing import Any, Callable, Coroutine, Dict, Generator, List, Optional
+from typing import Callable, Dict, Generator, List, Optional
 
 import justpy as jp
 from addict import Dict as AdDict
@@ -91,7 +91,7 @@ class Page(jp.QuasarPage):
             try:
                 if self.page_ready_generator is not None:
                     if isinstance(self.page_ready_generator, types.AsyncGeneratorType):
-                        await Automation(self.page_ready_generator.asend(PageEvent(msg.websocket)), context)
+                        await auto_context.update_after_await(self.page_ready_generator.asend(PageEvent(msg.websocket)), context)
                     elif isinstance(self.page_ready_generator, types.GeneratorType):
                         self.page_ready_generator.send(PageEvent(msg.websocket))
             except (StopIteration, StopAsyncIteration):
@@ -327,28 +327,3 @@ def create_favicon_routes() -> None:
     if globals.config.favicon:
         add_route(None, Route(f'/static/_favicon/{globals.config.favicon}',
                               lambda _: FileResponse(globals.config.favicon)))
-
-
-class Automation:
-
-    def __init__(self, coro: Coroutine, context):
-        self.coro = coro
-        self.context = context
-        self.context.lazy_update()
-
-    def __await__(self) -> Generator[Any, None, Any | None]:
-        coro_iter = self.coro.__await__()
-        iter_send, iter_throw = coro_iter.send, coro_iter.throw
-        send, message = iter_send, None
-        while True:
-            try:
-                signal = send(message)
-                self.context.lazy_update()
-            except StopIteration as err:
-                return err.value
-            else:
-                send = iter_send
-            try:
-                message = yield signal
-            except BaseException as err:
-                send, message = iter_throw, err