feat: on_app_start hook for app start and hot reload events

Author: maartenbreddels

solara/lab/__init__.py

@@ -1,7 +1,7 @@
 # isort: skip_file
 from .components import *  # noqa: F401, F403
 from .utils import cookies, headers  # noqa: F401, F403
-from ..lifecycle import on_kernel_start  # noqa: F401
+from ..lifecycle import on_kernel_start, on_app_start  # noqa: F401
 from ..tasks import task, use_task, Task, TaskResult  # noqa: F401, F403
 from ..toestand import computed  # noqa: F401
 

solara/lifecycle.py

@@ -11,7 +11,15 @@ class _on_kernel_callback_entry(NamedTuple):
     cleanup: Callable[[], None]
 
 
+class _on_app_start_callback_entry(NamedTuple):
+    callback: Callable[[], Optional[Callable[[], None]]]
+    callpoint: Optional[Path]
+    module: Optional[ModuleType]
+    cleanup: Callable[[], None]
+
+
 _on_kernel_start_callbacks: List[_on_kernel_callback_entry] = []
+_on_app_start_callbacks: List[_on_app_start_callback_entry] = []
 
 
 def _find_root_module_frame() -> Optional[FrameType]:
@@ -44,3 +52,53 @@ def cleanup():
     kce = _on_kernel_callback_entry(f, path, module, cleanup)
     _on_kernel_start_callbacks.append(kce)
     return cleanup
+
+
+def on_app_start(f: Callable[[], Optional[Callable[[], None]]]) -> Callable[[], None]:
+    """Run a function when your solara app starts and optionally run a cleanup function when hot reloading occurs.
+
+    `f` will be called on when you app is started using `solara run myapp.py`.
+    The (optional) function returned by `f` will be called when your app gets reloaded, which
+    happens [when you edit the app file and save it](/documentation/advanced/reference/reloading#reloading-of-python-files).
+
+    Note that the cleanup functions are called in reverse order with respect to the order in which they were registered
+    (e.g. the cleanup function of the last call to `on_app_start` will be called first).
+
+
+    If a cleanup function is not provided, you might as well not use `on_app_start` at all, and put your code directly in the module.
+
+    During hot reload, the callbacks that are added from scripts or modules that will be reloaded will be removed before the app is loaded
+    again. This can cause the order of the callbacks to be different than at first run.
+
+    ## Example
+
+    ```python
+    import solara
+    import solara.lab
+
+
+    @solara.lab.on_app_start
+    def app_start():
+        print("App started, initializing resources...")
+        def cleanup():
+            print("Cleaning up resources...")
+
+    ...
+    ```
+    """
+
+    root = _find_root_module_frame()
+    path: Optional[Path] = None
+    module: Optional[ModuleType] = None
+    if root is not None:
+        path_str = inspect.getsourcefile(root)
+        module = inspect.getmodule(root)
+        if path_str is not None:
+            path = Path(path_str)
+
+    def cleanup():
+        return _on_app_start_callbacks.remove(ace)
+
+    ace = _on_app_start_callback_entry(f, path, module, cleanup)
+    _on_app_start_callbacks.append(ace)
+    return cleanup

solara/server/app.py

@@ -9,7 +9,7 @@
 import warnings
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Callable, Dict, List, Optional, cast
 
 import ipywidgets as widgets
 import reacton
@@ -52,6 +52,7 @@ def __init__(self, name, default_app_name="Page"):
         if reload.reloader.on_change:
             raise RuntimeError("Previous reloader still had a on_change attached, no cleanup?")
         reload.reloader.on_change = self.on_file_change
+        self._on_app_close_callbacks: List[Callable[[], None]] = []
 
         self.app_name = default_app_name
         if ":" in self.fullname:
@@ -69,6 +70,7 @@ def __init__(self, name, default_app_name="Page"):
         if context is not None:
             raise RuntimeError(f"We should not have an existing Solara app context when running an app for the first time: {context}")
         dummy_kernel_context = kernel_context.create_dummy_context()
+
         with dummy_kernel_context:
             app = self._execute()
 
@@ -85,6 +87,11 @@ def __init__(self, name, default_app_name="Page"):
                 reload.reloader.root_path = package_root_path
         dummy_kernel_context.close()
 
+        for app_start_callback, *_ in solara.lifecycle._on_app_start_callbacks:
+            cleanup = app_start_callback()
+            if cleanup:
+                self._on_app_close_callbacks.append(cleanup)
+
     def _execute(self):
         logger.info("Executing %s", self.name)
         app = None
@@ -217,9 +224,16 @@ def run(self):
         if reload.reloader.requires_reload or self._first_execute_app is None:
             with thread_lock:
                 if reload.reloader.requires_reload or self._first_execute_app is None:
+                    required_reload = reload.reloader.requires_reload
                     self._first_execute_app = None
                     self._first_execute_app = self._execute()
                     print("Re-executed app", self.name)  # noqa
+                    if required_reload:
+                        # run after execute, which filled in the new _app_start callbacks
+                        for app_start_callback, *_ in solara.lifecycle._on_app_start_callbacks:
+                            cleanup = app_start_callback()
+                            if cleanup:
+                                self._on_app_close_callbacks.append(cleanup)
                     # We now ran the app again, might contain new imports
                     patch.patch_heavy_imports()
 
@@ -243,6 +257,9 @@ def reload(self):
         # if multiple files change in a short time, we want to do this
         # not concurrently. Even better would be to do a debounce?
         with thread_lock:
+            for cleanup in reversed(self._on_app_close_callbacks):
+                cleanup()
+            self._on_app_close_callbacks.clear()
             # TODO: clearing the type_counter is a bit of a hack
             # and we should introduce reload 'hooks', so there is
             # less interdependency between modules
@@ -275,6 +292,31 @@ def reload(self):
                     logger.info("reload: Removing on_kernel_start callback: %s (since it will be added when reloaded)", callback)
                     cleanup()
 
+            try:
+                for ac in solara.lifecycle._on_app_start_callbacks:
+                    callback, path, module, cleanup = ac
+                    will_reload = False
+                    if module is not None:
+                        module_name = module.__name__
+                        if module_name in reload.reloader.get_reload_module_names():
+                            will_reload = True
+                    elif path is not None:
+                        if str(path.resolve()).startswith(str(self.directory)):
+                            will_reload = True
+                        else:
+                            logger.warning(
+                                "script %s is not in the same directory as the app %s but is using on_app_start, "
+                                "this might lead to multiple entries, and might indicate a bug.",
+                                path,
+                                self.directory,
+                            )
+
+                    if will_reload:
+                        logger.info("reload: Removing on_app_start callback: %s (since it will be added when reloaded)", callback)
+                        cleanup()
+            except Exception as e:
+                logger.exception("Error while removing on_app_start callbacks: %s", e)
+
             context_values = list(kernel_context.contexts.values())
             # save states into the context so the hot reload will
             # keep the same state

solara/website/pages/documentation/api/utilities/on_app_start.py

@@ -0,0 +1,9 @@
+"""#on_app_start"""
+
+from solara.website.utils import apidoc
+import solara.lab
+from solara.website.components import NoPage
+
+title = "on_app_start"
+Page = NoPage
+__doc__ += apidoc(solara.lab.on_app_start)  # type: ignore

tests/unit/reload_test.py

@@ -11,6 +11,7 @@
 HERE = Path(__file__).parent
 
 kernel_start_path = HERE / "solara_test_apps" / "kernel_start.py"
+app_start_path = HERE / "solara_test_apps" / "app_start.py"
 
 
 @pytest.mark.parametrize("as_module", [False, True])
@@ -53,3 +54,47 @@ def test_callback_cleanup():
     assert test_callback_cleanup in [k.callback for k in solara.lifecycle._on_kernel_start_callbacks]
     cleanup()
     assert test_callback_cleanup not in [k.callback for k in solara.lifecycle._on_kernel_start_callbacks]
+
+
+@pytest.mark.parametrize("as_module", [False, True])
+def test_app_reload(tmpdir, kernel_context, extra_include_path, no_kernel_context, as_module):
+    target = Path(tmpdir) / "app_start.py"
+    shutil.copy(app_start_path, target)
+    with extra_include_path(str(tmpdir)):
+        on_app_start_callbacks = solara.lifecycle._on_app_start_callbacks.copy()
+        callbacks_start = [k.callback for k in solara.lifecycle._on_app_start_callbacks]
+        if as_module:
+            app = AppScript(f"{target.stem}")
+        else:
+            app = AppScript(f"{target}")
+        try:
+            app.run()
+            module = app.routes[0].module
+            module.started.assert_called_once()  # type: ignore
+            module.cleaned.assert_not_called()  # type: ignore
+            callback = module.app_start  # type: ignore
+            callbacks = [k.callback for k in solara.lifecycle._on_app_start_callbacks]
+            assert callbacks == [*callbacks_start, callback]
+            prev = callbacks.copy()
+            reload.reloader.reload_event_next.clear()
+            target.touch()
+            # wait for the event to trigger
+            reload.reloader.reload_event_next.wait()
+            module.started.assert_called_once()  # type: ignore
+            module.cleaned.assert_called_once()  # type: ignore
+            # we only 'rerun' after the first run
+            app.run()
+            module_reloaded = app.routes[0].module
+            module.started.assert_called_once()  # type: ignore
+            module.cleaned.assert_called_once()  # type: ignore
+            module_reloaded.started.assert_called_once()  # type: ignore
+            module_reloaded.cleaned.assert_not_called()  # type: ignore
+            assert module_reloaded is not module
+            callback = module_reloaded.app_start  # type: ignore
+            callbacks = [k[0] for k in solara.lifecycle._on_app_start_callbacks]
+            assert callbacks != prev
+            assert callbacks == [*callbacks_start, callback]
+        finally:
+            app.close()
+            solara.lifecycle._on_app_start_callbacks.clear()
+            solara.lifecycle._on_app_start_callbacks.extend(on_app_start_callbacks)