Skip to content

Commit 6f82e8f

Browse files
feat: on_app_start hook for app start and hot reload events
1 parent 99e0a00 commit 6f82e8f

File tree

5 files changed

+156
-2
lines changed

5 files changed

+156
-2
lines changed

solara/lab/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# isort: skip_file
22
from .components import * # noqa: F401, F403
33
from .utils import cookies, headers # noqa: F401, F403
4-
from ..lifecycle import on_kernel_start # noqa: F401
4+
from ..lifecycle import on_kernel_start, on_app_start # noqa: F401
55
from ..tasks import task, use_task, Task, TaskResult # noqa: F401, F403
66
from ..toestand import computed # noqa: F401
77

solara/lifecycle.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,15 @@ class _on_kernel_callback_entry(NamedTuple):
1111
cleanup: Callable[[], None]
1212

1313

14+
class _on_app_start_callback_entry(NamedTuple):
15+
callback: Callable[[], Optional[Callable[[], None]]]
16+
callpoint: Optional[Path]
17+
module: Optional[ModuleType]
18+
cleanup: Callable[[], None]
19+
20+
1421
_on_kernel_start_callbacks: List[_on_kernel_callback_entry] = []
22+
_on_app_start_callbacks: List[_on_app_start_callback_entry] = []
1523

1624

1725
def _find_root_module_frame() -> Optional[FrameType]:
@@ -44,3 +52,53 @@ def cleanup():
4452
kce = _on_kernel_callback_entry(f, path, module, cleanup)
4553
_on_kernel_start_callbacks.append(kce)
4654
return cleanup
55+
56+
57+
def on_app_start(f: Callable[[], Optional[Callable[[], None]]]) -> Callable[[], None]:
58+
"""Run a function when your solara app starts and optionally run a cleanup function when hot reloading occurs.
59+
60+
`f` will be called on when you app is started using `solara run myapp.py`.
61+
The (optional) function returned by `f` will be called when your app gets reloaded, which
62+
happens [when you edit the app file and save it](/documentation/advanced/reference/reloading#reloading-of-python-files).
63+
64+
Note that the cleanup functions are called in reverse order with respect to the order in which they were registered
65+
(e.g. the cleanup function of the last call to `on_app_start` will be called first).
66+
67+
68+
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.
69+
70+
During hot reload, the callbacks that are added from scripts or modules that will be reloaded will be removed before the app is loaded
71+
again. This can cause the order of the callbacks to be different than at first run.
72+
73+
## Example
74+
75+
```python
76+
import solara
77+
import solara.lab
78+
79+
80+
@solara.lab.on_app_start
81+
def app_start():
82+
print("App started, initializing resources...")
83+
def cleanup():
84+
print("Cleaning up resources...")
85+
86+
...
87+
```
88+
"""
89+
90+
root = _find_root_module_frame()
91+
path: Optional[Path] = None
92+
module: Optional[ModuleType] = None
93+
if root is not None:
94+
path_str = inspect.getsourcefile(root)
95+
module = inspect.getmodule(root)
96+
if path_str is not None:
97+
path = Path(path_str)
98+
99+
def cleanup():
100+
return _on_app_start_callbacks.remove(ace)
101+
102+
ace = _on_app_start_callback_entry(f, path, module, cleanup)
103+
_on_app_start_callbacks.append(ace)
104+
return cleanup

solara/server/app.py

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import warnings
1010
from enum import Enum
1111
from pathlib import Path
12-
from typing import Any, Dict, List, Optional, cast
12+
from typing import Any, Callable, Dict, List, Optional, cast
1313

1414
import ipywidgets as widgets
1515
import reacton
@@ -52,6 +52,7 @@ def __init__(self, name, default_app_name="Page"):
5252
if reload.reloader.on_change:
5353
raise RuntimeError("Previous reloader still had a on_change attached, no cleanup?")
5454
reload.reloader.on_change = self.on_file_change
55+
self._on_app_close_callbacks: List[Callable[[], None]] = []
5556

5657
self.app_name = default_app_name
5758
if ":" in self.fullname:
@@ -69,6 +70,7 @@ def __init__(self, name, default_app_name="Page"):
6970
if context is not None:
7071
raise RuntimeError(f"We should not have an existing Solara app context when running an app for the first time: {context}")
7172
dummy_kernel_context = kernel_context.create_dummy_context()
73+
7274
with dummy_kernel_context:
7375
app = self._execute()
7476

@@ -85,6 +87,11 @@ def __init__(self, name, default_app_name="Page"):
8587
reload.reloader.root_path = package_root_path
8688
dummy_kernel_context.close()
8789

90+
for app_start_callback, *_ in solara.lifecycle._on_app_start_callbacks:
91+
cleanup = app_start_callback()
92+
if cleanup:
93+
self._on_app_close_callbacks.append(cleanup)
94+
8895
def _execute(self):
8996
logger.info("Executing %s", self.name)
9097
app = None
@@ -217,9 +224,16 @@ def run(self):
217224
if reload.reloader.requires_reload or self._first_execute_app is None:
218225
with thread_lock:
219226
if reload.reloader.requires_reload or self._first_execute_app is None:
227+
required_reload = reload.reloader.requires_reload
220228
self._first_execute_app = None
221229
self._first_execute_app = self._execute()
222230
print("Re-executed app", self.name) # noqa
231+
if required_reload:
232+
# run after execute, which filled in the new _app_start callbacks
233+
for app_start_callback, *_ in solara.lifecycle._on_app_start_callbacks:
234+
cleanup = app_start_callback()
235+
if cleanup:
236+
self._on_app_close_callbacks.append(cleanup)
223237
# We now ran the app again, might contain new imports
224238
patch.patch_heavy_imports()
225239

@@ -243,6 +257,9 @@ def reload(self):
243257
# if multiple files change in a short time, we want to do this
244258
# not concurrently. Even better would be to do a debounce?
245259
with thread_lock:
260+
for cleanup in reversed(self._on_app_close_callbacks):
261+
cleanup()
262+
self._on_app_close_callbacks.clear()
246263
# TODO: clearing the type_counter is a bit of a hack
247264
# and we should introduce reload 'hooks', so there is
248265
# less interdependency between modules
@@ -275,6 +292,31 @@ def reload(self):
275292
logger.info("reload: Removing on_kernel_start callback: %s (since it will be added when reloaded)", callback)
276293
cleanup()
277294

295+
try:
296+
for ac in solara.lifecycle._on_app_start_callbacks:
297+
callback, path, module, cleanup = ac
298+
will_reload = False
299+
if module is not None:
300+
module_name = module.__name__
301+
if module_name in reload.reloader.get_reload_module_names():
302+
will_reload = True
303+
elif path is not None:
304+
if str(path.resolve()).startswith(str(self.directory)):
305+
will_reload = True
306+
else:
307+
logger.warning(
308+
"script %s is not in the same directory as the app %s but is using on_app_start, "
309+
"this might lead to multiple entries, and might indicate a bug.",
310+
path,
311+
self.directory,
312+
)
313+
314+
if will_reload:
315+
logger.info("reload: Removing on_app_start callback: %s (since it will be added when reloaded)", callback)
316+
cleanup()
317+
except Exception as e:
318+
logger.exception("Error while removing on_app_start callbacks: %s", e)
319+
278320
context_values = list(kernel_context.contexts.values())
279321
# save states into the context so the hot reload will
280322
# keep the same state
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
"""#on_app_start"""
2+
3+
from solara.website.utils import apidoc
4+
import solara.lab
5+
from solara.website.components import NoPage
6+
7+
title = "on_app_start"
8+
Page = NoPage
9+
__doc__ += apidoc(solara.lab.on_app_start) # type: ignore

tests/unit/reload_test.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
HERE = Path(__file__).parent
1212

1313
kernel_start_path = HERE / "solara_test_apps" / "kernel_start.py"
14+
app_start_path = HERE / "solara_test_apps" / "app_start.py"
1415

1516

1617
@pytest.mark.parametrize("as_module", [False, True])
@@ -53,3 +54,47 @@ def test_callback_cleanup():
5354
assert test_callback_cleanup in [k.callback for k in solara.lifecycle._on_kernel_start_callbacks]
5455
cleanup()
5556
assert test_callback_cleanup not in [k.callback for k in solara.lifecycle._on_kernel_start_callbacks]
57+
58+
59+
@pytest.mark.parametrize("as_module", [False, True])
60+
def test_app_reload(tmpdir, kernel_context, extra_include_path, no_kernel_context, as_module):
61+
target = Path(tmpdir) / "app_start.py"
62+
shutil.copy(app_start_path, target)
63+
with extra_include_path(str(tmpdir)):
64+
on_app_start_callbacks = solara.lifecycle._on_app_start_callbacks.copy()
65+
callbacks_start = [k.callback for k in solara.lifecycle._on_app_start_callbacks]
66+
if as_module:
67+
app = AppScript(f"{target.stem}")
68+
else:
69+
app = AppScript(f"{target}")
70+
try:
71+
app.run()
72+
module = app.routes[0].module
73+
module.started.assert_called_once() # type: ignore
74+
module.cleaned.assert_not_called() # type: ignore
75+
callback = module.app_start # type: ignore
76+
callbacks = [k.callback for k in solara.lifecycle._on_app_start_callbacks]
77+
assert callbacks == [*callbacks_start, callback]
78+
prev = callbacks.copy()
79+
reload.reloader.reload_event_next.clear()
80+
target.touch()
81+
# wait for the event to trigger
82+
reload.reloader.reload_event_next.wait()
83+
module.started.assert_called_once() # type: ignore
84+
module.cleaned.assert_called_once() # type: ignore
85+
# we only 'rerun' after the first run
86+
app.run()
87+
module_reloaded = app.routes[0].module
88+
module.started.assert_called_once() # type: ignore
89+
module.cleaned.assert_called_once() # type: ignore
90+
module_reloaded.started.assert_called_once() # type: ignore
91+
module_reloaded.cleaned.assert_not_called() # type: ignore
92+
assert module_reloaded is not module
93+
callback = module_reloaded.app_start # type: ignore
94+
callbacks = [k[0] for k in solara.lifecycle._on_app_start_callbacks]
95+
assert callbacks != prev
96+
assert callbacks == [*callbacks_start, callback]
97+
finally:
98+
app.close()
99+
solara.lifecycle._on_app_start_callbacks.clear()
100+
solara.lifecycle._on_app_start_callbacks.extend(on_app_start_callbacks)

0 commit comments

Comments
 (0)