diff --git a/mesop/features/__init__.py b/mesop/features/__init__.py index 5c063963e..c889f9baa 100644 --- a/mesop/features/__init__.py +++ b/mesop/features/__init__.py @@ -1 +1,2 @@ from .page import page as page +from .serve import serve as serve diff --git a/mesop/features/serve.py b/mesop/features/serve.py new file mode 100644 index 000000000..122f38c64 --- /dev/null +++ b/mesop/features/serve.py @@ -0,0 +1,17 @@ +from typing import Any, Callable + +from mesop.runtime import runtime + + +def serve( + *, + rule: str, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: + runtime().register_handler( + rule=rule, + handler=func, + ) + return func + + return decorator diff --git a/mesop/runtime/runtime.py b/mesop/runtime/runtime.py index 1f94b6cb7..f881fb843 100644 --- a/mesop/runtime/runtime.py +++ b/mesop/runtime/runtime.py @@ -1,6 +1,6 @@ from copy import deepcopy from dataclasses import dataclass -from typing import Any, Callable, Generator, Type, TypeVar, cast +from typing import Any, Callable, Dict, Generator, Tuple, Type, TypeVar, cast from flask import g, request @@ -49,6 +49,9 @@ def __init__(self): # clients polling whether to request a hot reload. self.hot_reload_counter = 0 self._path_to_page_config: dict[str, PageConfig] = {} + self._rule_to_handler: dict[ + str, Callable[[], Tuple[str, int, Dict[str, str]]] + ] = {} self.component_fns: set[Callable[..., Any]] = set() self.event_mappers: dict[Type[Any], Callable[[pb.UserEvent, Key], Any]] = {} self._state_classes: list[type[Any]] = [] @@ -131,9 +134,27 @@ def register_page(self, *, path: str, page_config: PageConfig) -> None: ) self._path_to_page_config[path] = page_config + def register_handler( + self, *, rule: str, handler: Callable[[], Tuple[str, int, Dict[str, str]]] + ): + if self._has_served_traffic: + raise MesopDeveloperException( + "Cannot register a handler after traffic has been served. You must register all handlers upon server startup before any traffic has been served. This prevents security issues." + ) + if rule in self._rule_to_handler: + raise MesopDeveloperException( + f"Handler for rule '{rule}' already registered." + ) + self._rule_to_handler[rule] = handler + def get_page_config(self, *, path: str) -> PageConfig | None: return self._path_to_page_config.get(path) + def get_rule_to_handlers( + self, + ) -> dict[str, Callable[[], Tuple[str, int, Dict[str, str]]]]: + return deepcopy(self._rule_to_handler) + def get_path_to_page_configs(self) -> dict[str, PageConfig]: return deepcopy(self._path_to_page_config) diff --git a/mesop/server/static_file_serving.py b/mesop/server/static_file_serving.py index 870931982..0b6421463 100644 --- a/mesop/server/static_file_serving.py +++ b/mesop/server/static_file_serving.py @@ -46,6 +46,10 @@ def configure_static_file_serving( disable_gzip_cache: bool = False, default_allowed_iframe_parents: str = "'self'", ): + # Register @me.serve handlers. + for rule, handler_func in runtime().get_rule_to_handlers().items(): + app.add_url_rule(rule, view_func=handler_func) + def get_path(path: str): safe_path = safe_join(static_file_runfiles_base, path) assert safe_path