diff --git a/sdk/python/packages/flet/src/flet/components/component.py b/sdk/python/packages/flet/src/flet/components/component.py index b151a7c1a1..62b4917dd1 100644 --- a/sdk/python/packages/flet/src/flet/components/component.py +++ b/sdk/python/packages/flet/src/flet/components/component.py @@ -341,7 +341,7 @@ def _run_render_effects(self): or hook.deps != hook.prev_deps ) if deps_changed: - if callable(hook.cleanup): + if callable(hook.cleanup) or callable(hook._returned_cleanup): self._schedule_effect(hook, is_cleanup=True) self._schedule_effect(hook, is_cleanup=False) @@ -354,7 +354,7 @@ def _run_unmount_effects(self): logger.debug("%s._run_unmount_effects()", self) for hook in self._state.hooks: # all effects are running on unmount - if isinstance(hook, EffectHook) and callable(hook.cleanup): + if isinstance(hook, EffectHook) and hook.has_lifecycle_work(): self._schedule_effect(hook, is_cleanup=True) self._state.hooks.clear() diff --git a/sdk/python/packages/flet/src/flet/components/hooks/use_effect.py b/sdk/python/packages/flet/src/flet/components/hooks/use_effect.py index 16dce1840a..fb88f6561f 100644 --- a/sdk/python/packages/flet/src/flet/components/hooks/use_effect.py +++ b/sdk/python/packages/flet/src/flet/components/hooks/use_effect.py @@ -20,6 +20,7 @@ class EffectHook(Hook): setup: Callable[[], Any | Awaitable[Any]] cleanup: Callable[[], Any | Awaitable[Any]] | None = None + _returned_cleanup: Callable[[], Any | Awaitable[Any]] | None = None deps: list[Any] | None = None prev_deps: list[Any] | None = None @@ -27,6 +28,13 @@ class EffectHook(Hook): _setup_task: asyncio.Task | None = None # last scheduled setup task _cleanup_task: asyncio.Task | None = None # last scheduled cleanup task + def has_lifecycle_work(self) -> bool: + return ( + callable(self.cleanup) + or callable(self._returned_cleanup) + or (self._setup_task is not None and not self._setup_task.done()) + ) + def cancel(self): """ Cancels currently scheduled setup and cleanup tasks, if still pending. diff --git a/sdk/python/packages/flet/src/flet/messaging/session.py b/sdk/python/packages/flet/src/flet/messaging/session.py index 30b396fbf6..22aaf920ad 100644 --- a/sdk/python/packages/flet/src/flet/messaging/session.py +++ b/sdk/python/packages/flet/src/flet/messaging/session.py @@ -735,14 +735,17 @@ async def __updates_scheduler(self): hook._setup_task = asyncio.create_task(hook.setup()) else: res = hook.setup() - if callable(res): - hook.cleanup = res - elif hook and hook.cleanup and is_cleanup: + hook._returned_cleanup = res if callable(res) else None + elif hook and is_cleanup: hook.cancel() - if inspect.iscoroutinefunction(hook.cleanup): - hook._cleanup_task = asyncio.create_task(hook.cleanup()) - else: - hook.cleanup() + cleanup = hook.cleanup or hook._returned_cleanup + hook._returned_cleanup = None + if callable(cleanup) and inspect.iscoroutinefunction( + cleanup + ): + hook._cleanup_task = asyncio.create_task(cleanup()) + elif callable(cleanup): + cleanup() except Exception as ex: tb = traceback.format_exc() self.error(f"Exception in effect: {ex}\n{tb}") diff --git a/sdk/python/packages/flet/tests/test_component_effects.py b/sdk/python/packages/flet/tests/test_component_effects.py index 8a09e43ca3..7151e1e2f5 100644 --- a/sdk/python/packages/flet/tests/test_component_effects.py +++ b/sdk/python/packages/flet/tests/test_component_effects.py @@ -1,3 +1,8 @@ +import asyncio + +import pytest + +import flet.components.hooks.use_effect as use_effect_module from flet.components.component import Component from flet.components.hooks.use_effect import EffectHook @@ -12,7 +17,7 @@ def test_render_effect_changed_deps_runs_cleanup_then_setup_once_each(): EffectHook( c, setup=lambda: None, - cleanup=lambda: None, + _returned_cleanup=lambda: None, deps=[2], prev_deps=[1], ) @@ -60,3 +65,56 @@ def test_will_unmount_releases_cached_render_references(): assert c._state.last_args == () assert c._state.last_kwargs == {} assert c._contexts == {} + + +def test_use_effect_retains_setup_returned_cleanup_across_renders(monkeypatch): + c = Component(fn=lambda: None, args=(), kwargs={}) + + def cleanup(): + return None + + def setup(): + return cleanup + + monkeypatch.setattr(use_effect_module, "current_component", lambda: c) + + use_effect_module.use_effect(setup, dependencies=[]) + hook = c._state.hooks[0] + hook._returned_cleanup = cleanup + + c._state.hook_cursor = 0 + use_effect_module.use_effect(setup, dependencies=[]) + + assert hook._returned_cleanup is cleanup + + +def test_unmount_schedules_returned_cleanup(): + c = Component(fn=lambda: None, args=(), kwargs={}) + c._state.mounted = True + calls: list[bool] = [] + c._schedule_effect = lambda hook, is_cleanup=False: calls.append(is_cleanup) # type: ignore[method-assign] + hook = EffectHook(c, setup=lambda: None, _returned_cleanup=lambda: None) + c._state.hooks = [hook] + + c._run_unmount_effects() + + assert calls == [True] + + +@pytest.mark.asyncio +async def test_unmount_schedules_cleanup_for_pending_setup_task(): + c = Component(fn=lambda: None, args=(), kwargs={}) + c._state.mounted = True + calls: list[bool] = [] + c._schedule_effect = lambda hook, is_cleanup=False: calls.append(is_cleanup) # type: ignore[method-assign] + task = asyncio.create_task(asyncio.sleep(60)) + hook = EffectHook(c, setup=lambda: None) + hook._setup_task = task + c._state.hooks = [hook] + + c._run_unmount_effects() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert calls == [True] diff --git a/sdk/python/packages/flet/tests/test_session_disconnect_buffering.py b/sdk/python/packages/flet/tests/test_session_disconnect_buffering.py index 2ef7110308..5d439f6f64 100644 --- a/sdk/python/packages/flet/tests/test_session_disconnect_buffering.py +++ b/sdk/python/packages/flet/tests/test_session_disconnect_buffering.py @@ -1,5 +1,8 @@ +import asyncio from datetime import datetime, timezone +import pytest + from flet.components.component import Component from flet.components.hooks.use_effect import EffectHook from flet.messaging.connection import Connection @@ -70,3 +73,56 @@ def test_attach_connection_restores_state_and_flushes_buffer(): assert session.expires_at is None assert session._Session__send_buffer == [] assert new_conn.messages == [buffered_message] + + +@pytest.mark.asyncio +async def test_cleanup_returned_by_effect_setup_runs_on_unmount(): + conn = Connection() + conn.pubsubhub = PubSubHub() + session = Session(conn) + component = Component(fn=lambda: None, args=(), kwargs={}) + component._state.mounted = True + cleanup_ran = asyncio.Event() + + def cleanup(): + cleanup_ran.set() + + hook = EffectHook(component, setup=lambda: None, _returned_cleanup=cleanup) + component._state.hooks = [hook] + component._schedule_effect = session.schedule_effect # type: ignore[method-assign] + component._run_unmount_effects() + session.start_updates_scheduler() + + await asyncio.wait_for(cleanup_ran.wait(), timeout=1) + session._Session__closed = True + session._Session__updates_ready.set() + await session._Session__updates_task + + +@pytest.mark.asyncio +async def test_unmount_cancels_pending_async_effect_setup(): + conn = Connection() + conn.pubsubhub = PubSubHub() + session = Session(conn) + component = Component(fn=lambda: None, args=(), kwargs={}) + component._state.mounted = True + started = asyncio.Event() + + async def setup(): + started.set() + await asyncio.Event().wait() + + task = asyncio.create_task(setup()) + await started.wait() + hook = EffectHook(component, setup=setup) + hook._setup_task = task + component._state.hooks = [hook] + component._schedule_effect = session.schedule_effect # type: ignore[method-assign] + component._run_unmount_effects() + session.start_updates_scheduler() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + session._Session__closed = True + session._Session__updates_ready.set() + await session._Session__updates_task