merge both subscribe methods

Author: falkoschindler

nicegui/event.py

@@ -12,6 +12,7 @@
 
 from . import background_tasks, core, helpers
 from .awaitable_response import AwaitableResponse
+from .client import Client
 from .context import context
 from .dataclasses import KWONLY_SLOTS
 from .logging import log
@@ -46,46 +47,44 @@ def __init__(self) -> None:
         self.callbacks: list[Callback[P]] = []
         self.instances.append(self)
 
-    def subscribe(self, callback: Callable[P, Any] | Callable[[], Any]) -> None:
+    def subscribe(self, callback: Callable[P, Any] | Callable[[], Any], *,
+                  unsubscribe_on_disconnect: bool | None = None) -> None:
         """Subscribe to the event.
 
-        Note that the callback should not be used to update UI since it would probably cause a memory leak.
-        Use the ``subscribe_ui`` method instead.
+        The ``unsubscribe_on_disconnect`` can be used to explicitly define
+        whether the callback should be automatically unsubscribed when the client disconnects.
+        By default, the callback is automatically unsubscribed if subscribed from within a UI context
+        to prevent memory leaks.
 
         :param callback: the callback which will be called when the event is fired
+        :param unsubscribe_on_disconnect: whether to unsubscribe the callback when the client disconnects
+            (default: ``None`` meaning the callback is automatically unsubscribed if subscribed from within a UI context)
         """
         frame = inspect.currentframe()
         assert frame is not None
         frame = frame.f_back
         assert frame is not None
-        self.callbacks.append(Callback(func=callback, filepath=frame.f_code.co_filename, line=frame.f_lineno))
-
-    def subscribe_ui(self, callback: Callable[P, Any] | Callable[[], Any]) -> None:
-        """Subscribe to the event and automatically unsubscribe when the client disconnects.
-
-        This method is particularly useful for UI callbacks since it cleans up unused references to the callback
-        which would otherwise cause a memory leak.
-
-        :param callback: the callback which will be called when the event is fired
-        """
-        self.subscribe(callback)
+        callback_ = Callback[P](func=callback, filepath=frame.f_code.co_filename, line=frame.f_lineno)
         try:
-            self.callbacks[-1].slot = weakref.ref(context.slot)
-        except RuntimeError as e:
-            raise RuntimeError('Calling `subscribe_ui` outside of a UI context is not supported.') from e
-        client = context.client
-
-        async def register_disconnect() -> None:
-            try:
-                await client.connected(timeout=10.0)
-                client.on_disconnect(lambda: self.unsubscribe(callback))
-            except TimeoutError:
-                log.warning('Could not register a disconnect handler for callback %s', callback)
-                self.unsubscribe(callback)
-        if core.loop and core.loop.is_running():
-            background_tasks.create(register_disconnect())
-        else:
-            core.app.on_startup(register_disconnect())
+            callback_.slot = weakref.ref(context.slot)
+            client: Client | None = context.client
+        except RuntimeError:
+            client = None
+        if callback_.slot is None and unsubscribe_on_disconnect is True:
+            raise RuntimeError('Calling `subscribe` with `unsubscribe_on_disconnect=True` outside of a UI context '
+                               'is not supported.')
+        if client is not None and unsubscribe_on_disconnect is not False:
+            async def register_disconnect() -> None:
+                try:
+                    await client.connected(timeout=10.0)
+                    client.on_disconnect(lambda: self.unsubscribe(callback))
+                except TimeoutError:
+                    log.warning('Could not register a disconnect handler for callback %s', callback)
+            if core.loop and core.loop.is_running():
+                background_tasks.create(register_disconnect())
+            else:
+                core.app.on_startup(register_disconnect())
+        self.callbacks.append(callback_)
 
     def unsubscribe(self, callback: Callable[P, Any] | Callable[[], Any]) -> None:
         """Unsubscribe a callback from the event.

tests/test_event.py

@@ -10,7 +10,7 @@ async def test_event(user: User):
     @ui.page('/')
     def page():
         ui.button('Click me', on_click=event.emit)
-        event.subscribe_ui(lambda: ui.notify('clicked'))
+        event.subscribe(lambda: ui.notify('clicked'))
 
     await user.open('/')
     user.find('Click me').click()
@@ -23,8 +23,8 @@ async def test_event_with_args(user: User):
     @ui.page('/')
     def page():
         ui.button('Click me', on_click=lambda: event.emit(42))
-        event.subscribe_ui(lambda: ui.notify('clicked'))
-        event.subscribe_ui(lambda x: ui.notify(f'{x = }'))
+        event.subscribe(lambda: ui.notify('clicked'))
+        event.subscribe(lambda x: ui.notify(f'{x = }'))
 
     await user.open('/')
     user.find('Click me').click()
@@ -39,7 +39,7 @@ async def test_event_with_async_handler(user: User):
     def page():
         ui.button('Click me', on_click=event.emit)
 
-        @event.subscribe_ui
+        @event.subscribe
         async def handler():
             await asyncio.sleep(0.1)
             ui.notify('clicked')
@@ -58,7 +58,7 @@ def page():
         nonlocal card
         ui.button('Click me', on_click=event.emit)
         with ui.card() as card:
-            event.subscribe_ui(lambda: ui.label('clicked'))
+            event.subscribe(lambda: ui.label('clicked'))
 
     await user.open('/')
     user.find('Click me').click()

website/documentation/content/event_documentation.py

@@ -9,32 +9,28 @@
     Events are a powerful tool distribute information between different parts of your code.
     The following demo shows how to define an event, subscribe a callback and emit it.
 
-    Note that there is also a `subscribe` method for non-UI callbacks.
-    But to run the handler in the current UI context and to automatically unsubscribe when the client disconnects,
-    use the `subscribe_ui` method.
-
     Handlers can be synchronous or asynchronous.
     They can also take arguments if the event contains arguments.
 ''')
 def events_demo():
     from nicegui import Event
 
     click = Event()
-    click.subscribe_ui(lambda: ui.notify('clicked'))
+    click.subscribe(lambda: ui.notify('clicked'))
 
     ui.button('Click me', on_click=click.emit)
 
 
 @doc.demo('Events with arguments', '''
     Events can also include arguments.
-    The callback can take use them, but also ignore them if they are not needed.
+    The callback can use them, but also ignore them if they are not needed.
 ''')
 def events_with_arguments():
     from nicegui import Event
 
     answer = Event[int]()
-    answer.subscribe_ui(lambda: ui.notify('Answer found!'))
-    answer.subscribe_ui(lambda x: ui.notify(f'{x = }'))
+    answer.subscribe(lambda: ui.notify('Answer found!'))
+    answer.subscribe(lambda x: ui.notify(f'{x = }'))
 
     ui.button('Answer', on_click=lambda: answer.emit(42))
 
@@ -51,7 +47,7 @@ def emitting_vs_calling_events():
 
     click = Event()
 
-    @click.subscribe_ui
+    @click.subscribe
     async def handler():
         n = ui.notification('Running...')
         await asyncio.sleep(1)