app.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989
  1. """The main Reflex app."""
  2. from __future__ import annotations
  3. import asyncio
  4. import contextlib
  5. import inspect
  6. import os
  7. from multiprocessing.pool import ThreadPool
  8. from typing import (
  9. Any,
  10. AsyncIterator,
  11. Callable,
  12. Coroutine,
  13. Dict,
  14. List,
  15. Optional,
  16. Set,
  17. Type,
  18. Union,
  19. )
  20. from fastapi import FastAPI, UploadFile
  21. from fastapi.middleware import cors
  22. from rich.progress import MofNCompleteColumn, Progress, TimeElapsedColumn
  23. from socketio import ASGIApp, AsyncNamespace, AsyncServer
  24. from starlette_admin.contrib.sqla.admin import Admin
  25. from starlette_admin.contrib.sqla.view import ModelView
  26. from reflex import constants
  27. from reflex.admin import AdminDash
  28. from reflex.base import Base
  29. from reflex.compiler import compiler
  30. from reflex.compiler import utils as compiler_utils
  31. from reflex.components import connection_modal
  32. from reflex.components.base.app_wrap import AppWrap
  33. from reflex.components.component import Component, ComponentStyle
  34. from reflex.components.layout.fragment import Fragment
  35. from reflex.components.navigation.client_side_routing import (
  36. Default404Page,
  37. wait_for_client_redirect,
  38. )
  39. from reflex.config import get_config
  40. from reflex.event import Event, EventHandler, EventSpec
  41. from reflex.middleware import HydrateMiddleware, Middleware
  42. from reflex.model import Model
  43. from reflex.page import (
  44. DECORATED_PAGES,
  45. )
  46. from reflex.route import (
  47. catchall_in_route,
  48. catchall_prefix,
  49. get_route_args,
  50. verify_route_validity,
  51. )
  52. from reflex.state import (
  53. DefaultState,
  54. State,
  55. StateManager,
  56. StateManagerMemory,
  57. StateUpdate,
  58. )
  59. from reflex.utils import console, format, prerequisites, types
  60. from reflex.vars import ImportVar
  61. # Define custom types.
  62. ComponentCallable = Callable[[], Component]
  63. Reducer = Callable[[Event], Coroutine[Any, Any, StateUpdate]]
  64. def default_overlay_component() -> Component:
  65. """Default overlay_component attribute for App.
  66. Returns:
  67. The default overlay_component, which is a connection_modal.
  68. """
  69. return connection_modal()
  70. class App(Base):
  71. """A Reflex application."""
  72. # A map from a page route to the component to render.
  73. pages: Dict[str, Component] = {}
  74. # A list of URLs to stylesheets to include in the app.
  75. stylesheets: List[str] = []
  76. # The backend API object.
  77. api: FastAPI = None # type: ignore
  78. # The Socket.IO AsyncServer.
  79. sio: Optional[AsyncServer] = None
  80. # The socket app.
  81. socket_app: Optional[ASGIApp] = None
  82. # The state class to use for the app.
  83. state: Type[State] = DefaultState
  84. # Class to manage many client states.
  85. state_manager: StateManager = StateManagerMemory(state=DefaultState)
  86. # The styling to apply to each component.
  87. style: ComponentStyle = {}
  88. # Middleware to add to the app.
  89. middleware: List[Middleware] = []
  90. # List of event handlers to trigger when a page loads.
  91. load_events: Dict[str, List[Union[EventHandler, EventSpec]]] = {}
  92. # Admin dashboard
  93. admin_dash: Optional[AdminDash] = None
  94. # The async server name space
  95. event_namespace: Optional[EventNamespace] = None
  96. # Components to add to the head of every page.
  97. head_components: List[Component] = []
  98. # A component that is present on every page.
  99. overlay_component: Optional[
  100. Union[Component, ComponentCallable]
  101. ] = default_overlay_component
  102. # Background tasks that are currently running
  103. background_tasks: Set[asyncio.Task] = set()
  104. # The radix theme for the entire app
  105. theme: Optional[Component] = None
  106. def __init__(self, *args, **kwargs):
  107. """Initialize the app.
  108. Args:
  109. *args: Args to initialize the app with.
  110. **kwargs: Kwargs to initialize the app with.
  111. Raises:
  112. ValueError: If the event namespace is not provided in the config.
  113. Also, if there are multiple client subclasses of rx.State(Subclasses of rx.State should consist
  114. of the DefaultState and the client app state).
  115. """
  116. if "connect_error_component" in kwargs:
  117. raise ValueError(
  118. "`connect_error_component` is deprecated, use `overlay_component` instead"
  119. )
  120. super().__init__(*args, **kwargs)
  121. state_subclasses = State.__subclasses__()
  122. inferred_state = state_subclasses[-1]
  123. is_testing_env = constants.PYTEST_CURRENT_TEST in os.environ
  124. # Special case to allow test cases have multiple subclasses of rx.State.
  125. if not is_testing_env:
  126. # Only the default state and the client state should be allowed as subclasses.
  127. if len(state_subclasses) > 2:
  128. raise ValueError(
  129. "rx.State has been subclassed multiple times. Only one subclass is allowed"
  130. )
  131. # verify that provided state is valid
  132. if self.state not in [DefaultState, inferred_state]:
  133. console.warn(
  134. f"Using substate ({self.state.__name__}) as root state in `rx.App` is currently not supported."
  135. f" Defaulting to root state: ({inferred_state.__name__})"
  136. )
  137. self.state = inferred_state
  138. # Get the config
  139. config = get_config()
  140. # Add middleware.
  141. self.middleware.append(HydrateMiddleware())
  142. # Set up the state manager.
  143. self.state_manager = StateManager.create(state=self.state)
  144. # Set up the API.
  145. self.api = FastAPI()
  146. self.add_cors()
  147. self.add_default_endpoints()
  148. if self.state is not DefaultState:
  149. # Set up the Socket.IO AsyncServer.
  150. self.sio = AsyncServer(
  151. async_mode="asgi",
  152. cors_allowed_origins="*"
  153. if config.cors_allowed_origins == ["*"]
  154. else config.cors_allowed_origins,
  155. cors_credentials=True,
  156. max_http_buffer_size=constants.POLLING_MAX_HTTP_BUFFER_SIZE,
  157. ping_interval=constants.Ping.INTERVAL,
  158. ping_timeout=constants.Ping.TIMEOUT,
  159. )
  160. # Create the socket app. Note event endpoint constant replaces the default 'socket.io' path.
  161. self.socket_app = ASGIApp(self.sio, socketio_path="")
  162. namespace = config.get_event_namespace()
  163. if not namespace:
  164. raise ValueError("event namespace must be provided in the config.")
  165. # Create the event namespace and attach the main app. Not related to any paths.
  166. self.event_namespace = EventNamespace(namespace, self)
  167. # Register the event namespace with the socket.
  168. self.sio.register_namespace(self.event_namespace)
  169. # Mount the socket app with the API.
  170. self.api.mount(str(constants.Endpoint.EVENT), self.socket_app)
  171. # Set up the admin dash.
  172. self.setup_admin_dash()
  173. # If a State is not used and no overlay_component is specified, do not render the connection modal
  174. if (
  175. self.state is DefaultState
  176. and self.overlay_component is default_overlay_component
  177. ):
  178. self.overlay_component = None
  179. def __repr__(self) -> str:
  180. """Get the string representation of the app.
  181. Returns:
  182. The string representation of the app.
  183. """
  184. return f"<App state={self.state.__name__}>"
  185. def __call__(self) -> FastAPI:
  186. """Run the backend api instance.
  187. Returns:
  188. The backend api.
  189. """
  190. return self.api
  191. def add_default_endpoints(self):
  192. """Add the default endpoints."""
  193. # To test the server.
  194. self.api.get(str(constants.Endpoint.PING))(ping)
  195. # To upload files.
  196. self.api.post(str(constants.Endpoint.UPLOAD))(upload(self))
  197. def add_cors(self):
  198. """Add CORS middleware to the app."""
  199. self.api.add_middleware(
  200. cors.CORSMiddleware,
  201. allow_credentials=True,
  202. allow_methods=["*"],
  203. allow_headers=["*"],
  204. allow_origins=["*"],
  205. )
  206. async def preprocess(self, state: State, event: Event) -> StateUpdate | None:
  207. """Preprocess the event.
  208. This is where middleware can modify the event before it is processed.
  209. Each middleware is called in the order it was added to the app.
  210. If a middleware returns an update, the event is not processed and the
  211. update is returned.
  212. Args:
  213. state: The state to preprocess.
  214. event: The event to preprocess.
  215. Returns:
  216. An optional state to return.
  217. """
  218. for middleware in self.middleware:
  219. if asyncio.iscoroutinefunction(middleware.preprocess):
  220. out = await middleware.preprocess(app=self, state=state, event=event) # type: ignore
  221. else:
  222. out = middleware.preprocess(app=self, state=state, event=event) # type: ignore
  223. if out is not None:
  224. return out # type: ignore
  225. async def postprocess(
  226. self, state: State, event: Event, update: StateUpdate
  227. ) -> StateUpdate:
  228. """Postprocess the event.
  229. This is where middleware can modify the delta after it is processed.
  230. Each middleware is called in the order it was added to the app.
  231. Args:
  232. state: The state to postprocess.
  233. event: The event to postprocess.
  234. update: The current state update.
  235. Returns:
  236. The state update to return.
  237. """
  238. for middleware in self.middleware:
  239. if asyncio.iscoroutinefunction(middleware.postprocess):
  240. out = await middleware.postprocess(
  241. app=self, state=state, event=event, update=update # type: ignore
  242. )
  243. else:
  244. out = middleware.postprocess(
  245. app=self, state=state, event=event, update=update # type: ignore
  246. )
  247. if out is not None:
  248. return out # type: ignore
  249. return update
  250. def add_middleware(self, middleware: Middleware, index: int | None = None):
  251. """Add middleware to the app.
  252. Args:
  253. middleware: The middleware to add.
  254. index: The index to add the middleware at.
  255. """
  256. if index is None:
  257. self.middleware.append(middleware)
  258. else:
  259. self.middleware.insert(index, middleware)
  260. @staticmethod
  261. def _generate_component(component: Component | ComponentCallable) -> Component:
  262. """Generate a component from a callable.
  263. Args:
  264. component: The component function to call or Component to return as-is.
  265. Returns:
  266. The generated component.
  267. Raises:
  268. TypeError: When an invalid component function is passed.
  269. """
  270. try:
  271. return component if isinstance(component, Component) else component()
  272. except TypeError as e:
  273. message = str(e)
  274. if "BaseVar" in message or "ComputedVar" in message:
  275. raise TypeError(
  276. "You may be trying to use an invalid Python function on a state var. "
  277. "When referencing a var inside your render code, only limited var operations are supported. "
  278. "See the var operation docs here: https://reflex.dev/docs/state/vars/#var-operations"
  279. ) from e
  280. raise e
  281. def add_page(
  282. self,
  283. component: Component | ComponentCallable,
  284. route: str | None = None,
  285. title: str = constants.DefaultPage.TITLE,
  286. description: str = constants.DefaultPage.DESCRIPTION,
  287. image: str = constants.DefaultPage.IMAGE,
  288. on_load: EventHandler
  289. | EventSpec
  290. | list[EventHandler | EventSpec]
  291. | None = None,
  292. meta: list[dict[str, str]] = constants.DefaultPage.META_LIST,
  293. script_tags: list[Component] | None = None,
  294. ):
  295. """Add a page to the app.
  296. If the component is a callable, by default the route is the name of the
  297. function. Otherwise, a route must be provided.
  298. Args:
  299. component: The component to display at the page.
  300. route: The route to display the component at.
  301. title: The title of the page.
  302. description: The description of the page.
  303. image: The image to display on the page.
  304. on_load: The event handler(s) that will be called each time the page load.
  305. meta: The metadata of the page.
  306. script_tags: List of script tags to be added to component
  307. """
  308. # If the route is not set, get it from the callable.
  309. if route is None:
  310. assert isinstance(
  311. component, Callable
  312. ), "Route must be set if component is not a callable."
  313. # Format the route.
  314. route = format.format_route(component.__name__)
  315. else:
  316. route = format.format_route(route, format_case=False)
  317. # Check if the route given is valid
  318. verify_route_validity(route)
  319. # Apply dynamic args to the route.
  320. self.state.setup_dynamic_args(get_route_args(route))
  321. # Generate the component if it is a callable.
  322. component = self._generate_component(component)
  323. # Wrap the component in a fragment with optional overlay.
  324. if self.overlay_component is not None:
  325. component = Fragment.create(
  326. self._generate_component(self.overlay_component),
  327. component,
  328. )
  329. else:
  330. component = Fragment.create(component)
  331. # Add meta information to the component.
  332. compiler_utils.add_meta(
  333. component,
  334. title=title,
  335. image=image,
  336. description=description,
  337. meta=meta,
  338. )
  339. # Add script tags if given
  340. if script_tags:
  341. console.deprecate(
  342. feature_name="Passing script tags to add_page",
  343. reason="Add script components as children to the page component instead",
  344. deprecation_version="0.2.9",
  345. removal_version="0.3.1",
  346. )
  347. component.children.extend(script_tags)
  348. # Add the page.
  349. self._check_routes_conflict(route)
  350. self.pages[route] = component
  351. # Add the load events.
  352. if on_load:
  353. if not isinstance(on_load, list):
  354. on_load = [on_load]
  355. self.load_events[route] = on_load
  356. def get_load_events(self, route: str) -> list[EventHandler | EventSpec]:
  357. """Get the load events for a route.
  358. Args:
  359. route: The route to get the load events for.
  360. Returns:
  361. The load events for the route.
  362. """
  363. route = route.lstrip("/")
  364. if route == "":
  365. route = constants.PageNames.INDEX_ROUTE
  366. return self.load_events.get(route, [])
  367. def _check_routes_conflict(self, new_route: str):
  368. """Verify if there is any conflict between the new route and any existing route.
  369. Based on conflicts that NextJS would throw if not intercepted.
  370. Raises:
  371. ValueError: exception showing which conflict exist with the route to be added
  372. Args:
  373. new_route: the route being newly added.
  374. """
  375. newroute_catchall = catchall_in_route(new_route)
  376. if not newroute_catchall:
  377. return
  378. for route in self.pages:
  379. route = "" if route == "index" else route
  380. if new_route.startswith(f"{route}/[[..."):
  381. raise ValueError(
  382. f"You cannot define a route with the same specificity as a optional catch-all route ('{route}' and '{new_route}')"
  383. )
  384. route_catchall = catchall_in_route(route)
  385. if (
  386. route_catchall
  387. and newroute_catchall
  388. and catchall_prefix(route) == catchall_prefix(new_route)
  389. ):
  390. raise ValueError(
  391. f"You cannot use multiple catchall for the same dynamic route ({route} !== {new_route})"
  392. )
  393. def add_custom_404_page(
  394. self,
  395. component: Component | ComponentCallable | None = None,
  396. title: str = constants.Page404.TITLE,
  397. image: str = constants.Page404.IMAGE,
  398. description: str = constants.Page404.DESCRIPTION,
  399. on_load: EventHandler
  400. | EventSpec
  401. | list[EventHandler | EventSpec]
  402. | None = None,
  403. meta: list[dict[str, str]] = constants.DefaultPage.META_LIST,
  404. ):
  405. """Define a custom 404 page for any url having no match.
  406. If there is no page defined on 'index' route, add the 404 page to it.
  407. If there is no global catchall defined, add the 404 page with a catchall
  408. Args:
  409. component: The component to display at the page.
  410. title: The title of the page.
  411. description: The description of the page.
  412. image: The image to display on the page.
  413. on_load: The event handler(s) that will be called each time the page load.
  414. meta: The metadata of the page.
  415. """
  416. if component is None:
  417. component = Default404Page.create()
  418. self.add_page(
  419. component=wait_for_client_redirect(self._generate_component(component)),
  420. route=constants.Page404.SLUG,
  421. title=title or constants.Page404.TITLE,
  422. image=image or constants.Page404.IMAGE,
  423. description=description or constants.Page404.DESCRIPTION,
  424. on_load=on_load,
  425. meta=meta,
  426. )
  427. def setup_admin_dash(self):
  428. """Setup the admin dash."""
  429. # Get the admin dash.
  430. admin_dash = self.admin_dash
  431. if admin_dash and admin_dash.models:
  432. # Build the admin dashboard
  433. admin = (
  434. admin_dash.admin
  435. if admin_dash.admin
  436. else Admin(
  437. engine=Model.get_db_engine(),
  438. title="Reflex Admin Dashboard",
  439. logo_url="https://reflex.dev/Reflex.svg",
  440. )
  441. )
  442. for model in admin_dash.models:
  443. view = admin_dash.view_overrides.get(model, ModelView)
  444. admin.add_view(view(model))
  445. admin.mount_to(self.api)
  446. def get_frontend_packages(self, imports: Dict[str, set[ImportVar]]):
  447. """Gets the frontend packages to be installed and filters out the unnecessary ones.
  448. Args:
  449. imports: A dictionary containing the imports used in the current page.
  450. Example:
  451. >>> get_frontend_packages({"react": "16.14.0", "react-dom": "16.14.0"})
  452. """
  453. page_imports = {
  454. i
  455. for i, tags in imports.items()
  456. if i
  457. not in [
  458. *compiler.DEFAULT_IMPORTS.keys(),
  459. *constants.PackageJson.DEPENDENCIES.keys(),
  460. *constants.PackageJson.DEV_DEPENDENCIES.keys(),
  461. ]
  462. and not any(i.startswith(prefix) for prefix in ["/", ".", "next/"])
  463. and i != ""
  464. and any(tag.install for tag in tags)
  465. }
  466. frontend_packages = get_config().frontend_packages
  467. _frontend_packages = []
  468. for package in frontend_packages:
  469. if package in (get_config().tailwind or {}).get("plugins", []): # type: ignore
  470. console.warn(
  471. f"Tailwind packages are inferred from 'plugins', remove `{package}` from `frontend_packages`"
  472. )
  473. continue
  474. if package in page_imports:
  475. console.warn(
  476. f"React packages and their dependencies are inferred from Component.library and Component.lib_dependencies, remove `{package}` from `frontend_packages`"
  477. )
  478. continue
  479. _frontend_packages.append(package)
  480. page_imports.update(_frontend_packages)
  481. prerequisites.install_frontend_packages(page_imports)
  482. def _app_root(self, app_wrappers):
  483. order = sorted(app_wrappers, key=lambda k: k[0], reverse=True)
  484. root = parent = app_wrappers[order[0]]
  485. for key in order[1:]:
  486. child = app_wrappers[key]
  487. parent.children.append(child)
  488. parent = child
  489. return root
  490. def compile(self):
  491. """Compile the app and output it to the pages folder."""
  492. if os.environ.get(constants.SKIP_COMPILE_ENV_VAR) == "yes":
  493. return
  494. # Create a progress bar.
  495. progress = Progress(
  496. *Progress.get_default_columns()[:-1],
  497. MofNCompleteColumn(),
  498. TimeElapsedColumn(),
  499. )
  500. for render, kwargs in DECORATED_PAGES:
  501. self.add_page(render, **kwargs)
  502. # Render a default 404 page if the user didn't supply one
  503. if constants.Page404.SLUG not in self.pages:
  504. self.add_custom_404_page()
  505. task = progress.add_task("Compiling: ", total=len(self.pages))
  506. # TODO: include all work done in progress indicator, not just self.pages
  507. # Get the env mode.
  508. config = get_config()
  509. # Store the compile results.
  510. compile_results = []
  511. # Compile the pages in parallel.
  512. custom_components = set()
  513. # TODO Anecdotally, processes=2 works 10% faster (cpu_count=12)
  514. thread_pool = ThreadPool()
  515. all_imports = {}
  516. page_futures = []
  517. app_wrappers: Dict[tuple[int, str], Component] = {
  518. # Default app wrap component renders {children}
  519. (0, "AppWrap"): AppWrap.create()
  520. }
  521. if self.theme is not None:
  522. # If a theme component was provided, wrap the app with it
  523. app_wrappers[(20, "Theme")] = self.theme
  524. with progress:
  525. for route, component in self.pages.items():
  526. # TODO: this progress does not reflect actual threaded task completion
  527. progress.advance(task)
  528. component.add_style(self.style)
  529. page_futures.append(
  530. thread_pool.apply_async(
  531. compiler.compile_page,
  532. args=(
  533. route,
  534. component,
  535. self.state,
  536. ),
  537. )
  538. )
  539. # add component.get_imports() to all_imports
  540. all_imports.update(component.get_imports())
  541. # add the app wrappers from this component
  542. app_wrappers.update(component.get_app_wrap_components())
  543. # Add the custom components from the page to the set.
  544. custom_components |= component.get_custom_components()
  545. thread_pool.close()
  546. thread_pool.join()
  547. # Compile the app wrapper.
  548. app_root = self._app_root(app_wrappers=app_wrappers)
  549. all_imports.update(app_root.get_imports())
  550. compile_results.append(compiler.compile_app(app_root))
  551. # Get the compiled pages.
  552. compile_results.extend(result.get() for result in page_futures)
  553. # TODO the compile tasks below may also benefit from parallelization too
  554. # Compile the custom components.
  555. compile_results.append(compiler.compile_components(custom_components))
  556. # Iterate through all the custom components and add their imports to the all_imports
  557. for component in custom_components:
  558. all_imports.update(component.get_imports())
  559. # Compile the root stylesheet with base styles.
  560. compile_results.append(compiler.compile_root_stylesheet(self.stylesheets))
  561. # Compile the root document.
  562. compile_results.append(compiler.compile_document_root(self.head_components))
  563. # Compile the theme.
  564. compile_results.append(compiler.compile_theme(style=self.style))
  565. # Compile the contexts.
  566. compile_results.append(compiler.compile_contexts(self.state))
  567. # Compile the Tailwind config.
  568. if config.tailwind is not None:
  569. config.tailwind["content"] = config.tailwind.get(
  570. "content", constants.Tailwind.CONTENT
  571. )
  572. compile_results.append(compiler.compile_tailwind(config.tailwind))
  573. # Empty the .web pages directory
  574. compiler.purge_web_pages_dir()
  575. # install frontend packages
  576. self.get_frontend_packages(all_imports)
  577. # Write the pages at the end to trigger the NextJS hot reload only once.
  578. thread_pool = ThreadPool()
  579. for output_path, code in compile_results:
  580. thread_pool.apply_async(compiler_utils.write_page, args=(output_path, code))
  581. thread_pool.close()
  582. thread_pool.join()
  583. @contextlib.asynccontextmanager
  584. async def modify_state(self, token: str) -> AsyncIterator[State]:
  585. """Modify the state out of band.
  586. Args:
  587. token: The token to modify the state for.
  588. Yields:
  589. The state to modify.
  590. Raises:
  591. RuntimeError: If the app has not been initialized yet.
  592. """
  593. if self.event_namespace is None:
  594. raise RuntimeError("App has not been initialized yet.")
  595. # Get exclusive access to the state.
  596. async with self.state_manager.modify_state(token) as state:
  597. # No other event handler can modify the state while in this context.
  598. yield state
  599. delta = state.get_delta()
  600. if delta:
  601. # When the state is modified reset dirty status and emit the delta to the frontend.
  602. state._clean()
  603. await self.event_namespace.emit_update(
  604. update=StateUpdate(delta=delta),
  605. sid=state.get_sid(),
  606. )
  607. def _process_background(self, state: State, event: Event) -> asyncio.Task | None:
  608. """Process an event in the background and emit updates as they arrive.
  609. Args:
  610. state: The state to process the event for.
  611. event: The event to process.
  612. Returns:
  613. Task if the event was backgroundable, otherwise None
  614. """
  615. substate, handler = state._get_event_handler(event)
  616. if not handler.is_background:
  617. return None
  618. async def _coro():
  619. """Coroutine to process the event and emit updates inside an asyncio.Task.
  620. Raises:
  621. RuntimeError: If the app has not been initialized yet.
  622. """
  623. if self.event_namespace is None:
  624. raise RuntimeError("App has not been initialized yet.")
  625. # Process the event.
  626. async for update in state._process_event(
  627. handler=handler, state=substate, payload=event.payload
  628. ):
  629. # Postprocess the event.
  630. update = await self.postprocess(state, event, update)
  631. # Send the update to the client.
  632. await self.event_namespace.emit_update(
  633. update=update,
  634. sid=state.get_sid(),
  635. )
  636. task = asyncio.create_task(_coro())
  637. self.background_tasks.add(task)
  638. # Clean up task from background_tasks set when complete.
  639. task.add_done_callback(self.background_tasks.discard)
  640. return task
  641. async def process(
  642. app: App, event: Event, sid: str, headers: Dict, client_ip: str
  643. ) -> AsyncIterator[StateUpdate]:
  644. """Process an event.
  645. Args:
  646. app: The app to process the event for.
  647. event: The event to process.
  648. sid: The Socket.IO session id.
  649. headers: The client headers.
  650. client_ip: The client_ip.
  651. Yields:
  652. The state updates after processing the event.
  653. """
  654. # Add request data to the state.
  655. router_data = event.router_data
  656. router_data.update(
  657. {
  658. constants.RouteVar.QUERY: format.format_query_params(event.router_data),
  659. constants.RouteVar.CLIENT_TOKEN: event.token,
  660. constants.RouteVar.SESSION_ID: sid,
  661. constants.RouteVar.HEADERS: headers,
  662. constants.RouteVar.CLIENT_IP: client_ip,
  663. }
  664. )
  665. # Get the state for the session exclusively.
  666. async with app.state_manager.modify_state(event.token) as state:
  667. # re-assign only when the value is different
  668. if state.router_data != router_data:
  669. # assignment will recurse into substates and force recalculation of
  670. # dependent ComputedVar (dynamic route variables)
  671. state.router_data = router_data
  672. # Preprocess the event.
  673. update = await app.preprocess(state, event)
  674. # If there was an update, yield it.
  675. if update is not None:
  676. yield update
  677. # Only process the event if there is no update.
  678. else:
  679. if app._process_background(state, event) is not None:
  680. # `final=True` allows the frontend send more events immediately.
  681. yield StateUpdate(final=True)
  682. return
  683. # Process the event synchronously.
  684. async for update in state._process(event):
  685. # Postprocess the event.
  686. update = await app.postprocess(state, event, update)
  687. # Yield the update.
  688. yield update
  689. async def ping() -> str:
  690. """Test API endpoint.
  691. Returns:
  692. The response.
  693. """
  694. return "pong"
  695. def upload(app: App):
  696. """Upload a file.
  697. Args:
  698. app: The app to upload the file for.
  699. Returns:
  700. The upload function.
  701. """
  702. async def upload_file(files: List[UploadFile]):
  703. """Upload a file.
  704. Args:
  705. files: The file(s) to upload.
  706. Raises:
  707. ValueError: if there are no args with supported annotation.
  708. """
  709. assert files[0].filename is not None
  710. token, handler = files[0].filename.split(":")[:2]
  711. for file in files:
  712. assert file.filename is not None
  713. file.filename = file.filename.split(":")[-1]
  714. # Get the state for the session.
  715. async with app.state_manager.modify_state(token) as state:
  716. # get the current session ID
  717. sid = state.get_sid()
  718. # get the current state(parent state/substate)
  719. path = handler.split(".")[:-1]
  720. current_state = state.get_substate(path)
  721. handler_upload_param = ()
  722. # get handler function
  723. func = getattr(current_state, handler.split(".")[-1])
  724. # check if there exists any handler args with annotation, List[UploadFile]
  725. for k, v in inspect.getfullargspec(
  726. func.fn if isinstance(func, EventHandler) else func
  727. ).annotations.items():
  728. if types.is_generic_alias(v) and types._issubclass(
  729. v.__args__[0], UploadFile
  730. ):
  731. handler_upload_param = (k, v)
  732. break
  733. if not handler_upload_param:
  734. raise ValueError(
  735. f"`{handler}` handler should have a parameter annotated as List["
  736. f"rx.UploadFile]"
  737. )
  738. event = Event(
  739. token=token,
  740. name=handler,
  741. payload={handler_upload_param[0]: files},
  742. )
  743. async for update in state._process(event):
  744. # Postprocess the event.
  745. update = await app.postprocess(state, event, update)
  746. # Send update to client
  747. await app.event_namespace.emit_update( # type: ignore
  748. update=update,
  749. sid=sid,
  750. )
  751. return upload_file
  752. class EventNamespace(AsyncNamespace):
  753. """The event namespace."""
  754. # The application object.
  755. app: App
  756. def __init__(self, namespace: str, app: App):
  757. """Initialize the event namespace.
  758. Args:
  759. namespace: The namespace.
  760. app: The application object.
  761. """
  762. super().__init__(namespace)
  763. self.app = app
  764. def on_connect(self, sid, environ):
  765. """Event for when the websocket is connected.
  766. Args:
  767. sid: The Socket.IO session id.
  768. environ: The request information, including HTTP headers.
  769. """
  770. pass
  771. def on_disconnect(self, sid):
  772. """Event for when the websocket disconnects.
  773. Args:
  774. sid: The Socket.IO session id.
  775. """
  776. pass
  777. async def emit_update(self, update: StateUpdate, sid: str) -> None:
  778. """Emit an update to the client.
  779. Args:
  780. update: The state update to send.
  781. sid: The Socket.IO session id.
  782. """
  783. # Creating a task prevents the update from being blocked behind other coroutines.
  784. await asyncio.create_task(
  785. self.emit(str(constants.SocketEvent.EVENT), update.json(), to=sid)
  786. )
  787. async def on_event(self, sid, data):
  788. """Event for receiving front-end websocket events.
  789. Args:
  790. sid: The Socket.IO session id.
  791. data: The event data.
  792. """
  793. # Get the event.
  794. event = Event.parse_raw(data)
  795. # Get the event environment.
  796. assert self.app.sio is not None
  797. environ = self.app.sio.get_environ(sid, self.namespace)
  798. assert environ is not None
  799. # Get the client headers.
  800. headers = {
  801. k.decode("utf-8"): v.decode("utf-8")
  802. for (k, v) in environ["asgi.scope"]["headers"]
  803. }
  804. # Get the client IP
  805. client_ip = environ["REMOTE_ADDR"]
  806. # Process the events.
  807. async for update in process(self.app, event, sid, headers, client_ip):
  808. # Emit the update from processing the event.
  809. await self.emit_update(update=update, sid=sid)
  810. async def on_ping(self, sid):
  811. """Event for testing the API endpoint.
  812. Args:
  813. sid: The Socket.IO session id.
  814. """
  815. # Emit the test event.
  816. await self.emit(str(constants.SocketEvent.PING), "pong", to=sid)