state.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994
  1. """Define the pynecone state specification."""
  2. from __future__ import annotations
  3. import asyncio
  4. import copy
  5. import functools
  6. import inspect
  7. import traceback
  8. from abc import ABC
  9. from collections import defaultdict
  10. from typing import (
  11. Any,
  12. AsyncIterator,
  13. Callable,
  14. ClassVar,
  15. Dict,
  16. List,
  17. Optional,
  18. Sequence,
  19. Set,
  20. Tuple,
  21. Type,
  22. Union,
  23. )
  24. import cloudpickle
  25. import pydantic
  26. from redis import Redis
  27. from pynecone import constants
  28. from pynecone.base import Base
  29. from pynecone.event import Event, EventHandler, EventSpec, fix_events, window_alert
  30. from pynecone.utils import format, prerequisites, types
  31. from pynecone.vars import BaseVar, ComputedVar, PCDict, PCList, Var
  32. Delta = Dict[str, Any]
  33. class State(Base, ABC, extra=pydantic.Extra.allow):
  34. """The state of the app."""
  35. # A map from the var name to the var.
  36. vars: ClassVar[Dict[str, Var]] = {}
  37. # The base vars of the class.
  38. base_vars: ClassVar[Dict[str, BaseVar]] = {}
  39. # The computed vars of the class.
  40. computed_vars: ClassVar[Dict[str, ComputedVar]] = {}
  41. # Vars inherited by the parent state.
  42. inherited_vars: ClassVar[Dict[str, Var]] = {}
  43. # Backend vars that are never sent to the client.
  44. backend_vars: ClassVar[Dict[str, Any]] = {}
  45. # Backend vars inherited
  46. inherited_backend_vars: ClassVar[Dict[str, Any]] = {}
  47. # The event handlers.
  48. event_handlers: ClassVar[Dict[str, EventHandler]] = {}
  49. # The parent state.
  50. parent_state: Optional[State] = None
  51. # The substates of the state.
  52. substates: Dict[str, State] = {}
  53. # The set of dirty vars.
  54. dirty_vars: Set[str] = set()
  55. # The set of dirty substates.
  56. dirty_substates: Set[str] = set()
  57. # The routing path that triggered the state
  58. router_data: Dict[str, Any] = {}
  59. # Mapping of var name to set of computed variables that depend on it
  60. computed_var_dependencies: Dict[str, Set[str]] = {}
  61. # Mapping of var name to set of substates that depend on it
  62. substate_var_dependencies: Dict[str, Set[str]] = {}
  63. # Per-instance copy of backend variable values
  64. _backend_vars: Dict[str, Any] = {}
  65. def __init__(self, *args, parent_state: Optional[State] = None, **kwargs):
  66. """Initialize the state.
  67. Args:
  68. *args: The args to pass to the Pydantic init method.
  69. parent_state: The parent state.
  70. **kwargs: The kwargs to pass to the Pydantic init method.
  71. """
  72. kwargs["parent_state"] = parent_state
  73. super().__init__(*args, **kwargs)
  74. # initialize per-instance var dependency tracking
  75. self.computed_var_dependencies = defaultdict(set)
  76. self.substate_var_dependencies = defaultdict(set)
  77. # Setup the substates.
  78. for substate in self.get_substates():
  79. self.substates[substate.get_name()] = substate(parent_state=self)
  80. # Convert the event handlers to functions.
  81. for name, event_handler in self.event_handlers.items():
  82. fn = functools.partial(event_handler.fn, self)
  83. fn.__module__ = event_handler.fn.__module__ # type: ignore
  84. fn.__qualname__ = event_handler.fn.__qualname__ # type: ignore
  85. setattr(self, name, fn)
  86. # Initialize computed vars dependencies.
  87. inherited_vars = set(self.inherited_vars).union(
  88. set(self.inherited_backend_vars),
  89. )
  90. for cvar_name, cvar in self.computed_vars.items():
  91. # Add the dependencies.
  92. for var in cvar.deps(objclass=type(self)):
  93. self.computed_var_dependencies[var].add(cvar_name)
  94. if var in inherited_vars:
  95. # track that this substate depends on its parent for this var
  96. state_name = self.get_name()
  97. parent_state = self.parent_state
  98. while parent_state is not None and var in parent_state.vars:
  99. parent_state.substate_var_dependencies[var].add(state_name)
  100. state_name, parent_state = (
  101. parent_state.get_name(),
  102. parent_state.parent_state,
  103. )
  104. # Initialize the mutable fields.
  105. self._init_mutable_fields()
  106. # Create a fresh copy of the backend variables for this instance
  107. self._backend_vars = copy.deepcopy(self.backend_vars)
  108. def _init_mutable_fields(self):
  109. """Initialize mutable fields.
  110. So that mutation to them can be detected by the app:
  111. * list
  112. """
  113. for field in self.base_vars.values():
  114. value = getattr(self, field.name)
  115. value_in_pc_data = _convert_mutable_datatypes(
  116. value, self._reassign_field, field.name
  117. )
  118. if types._issubclass(field.type_, Union[List, Dict]):
  119. setattr(self, field.name, value_in_pc_data)
  120. self.clean()
  121. def _reassign_field(self, field_name: str):
  122. """Reassign the given field.
  123. Primarily for mutation in fields of mutable data types.
  124. Args:
  125. field_name: The name of the field we want to reassign
  126. """
  127. setattr(
  128. self,
  129. field_name,
  130. getattr(self, field_name),
  131. )
  132. def __repr__(self) -> str:
  133. """Get the string representation of the state.
  134. Returns:
  135. The string representation of the state.
  136. """
  137. return f"{self.__class__.__name__}({self.dict()})"
  138. @classmethod
  139. def __init_subclass__(cls, **kwargs):
  140. """Do some magic for the subclass initialization.
  141. Args:
  142. **kwargs: The kwargs to pass to the pydantic init_subclass method.
  143. """
  144. super().__init_subclass__(**kwargs)
  145. # Get the parent vars.
  146. parent_state = cls.get_parent_state()
  147. if parent_state is not None:
  148. cls.inherited_vars = parent_state.vars
  149. cls.inherited_backend_vars = parent_state.backend_vars
  150. cls.new_backend_vars = {
  151. name: value
  152. for name, value in cls.__dict__.items()
  153. if types.is_backend_variable(name)
  154. and name not in cls.inherited_backend_vars
  155. }
  156. cls.backend_vars = {**cls.inherited_backend_vars, **cls.new_backend_vars}
  157. # Set the base and computed vars.
  158. cls.base_vars = {
  159. f.name: BaseVar(name=f.name, type_=f.outer_type_).set_state(cls)
  160. for f in cls.get_fields().values()
  161. if f.name not in cls.get_skip_vars()
  162. }
  163. cls.computed_vars = {
  164. v.name: v.set_state(cls)
  165. for v in cls.__dict__.values()
  166. if isinstance(v, ComputedVar)
  167. }
  168. cls.vars = {
  169. **cls.inherited_vars,
  170. **cls.base_vars,
  171. **cls.computed_vars,
  172. }
  173. cls.event_handlers = {}
  174. # Setup the base vars at the class level.
  175. for prop in cls.base_vars.values():
  176. cls._init_var(prop)
  177. # Set up the event handlers.
  178. events = {
  179. name: fn
  180. for name, fn in cls.__dict__.items()
  181. if not name.startswith("_")
  182. and isinstance(fn, Callable)
  183. and not isinstance(fn, EventHandler)
  184. }
  185. for name, fn in events.items():
  186. handler = EventHandler(fn=fn)
  187. cls.event_handlers[name] = handler
  188. setattr(cls, name, handler)
  189. @classmethod
  190. def get_skip_vars(cls) -> Set[str]:
  191. """Get the vars to skip when serializing.
  192. Returns:
  193. The vars to skip when serializing.
  194. """
  195. return set(cls.inherited_vars) | {
  196. "parent_state",
  197. "substates",
  198. "dirty_vars",
  199. "dirty_substates",
  200. "router_data",
  201. "computed_var_dependencies",
  202. "substate_var_dependencies",
  203. "_backend_vars",
  204. }
  205. @classmethod
  206. @functools.lru_cache()
  207. def get_parent_state(cls) -> Optional[Type[State]]:
  208. """Get the parent state.
  209. Returns:
  210. The parent state.
  211. """
  212. parent_states = [
  213. base
  214. for base in cls.__bases__
  215. if types._issubclass(base, State) and base is not State
  216. ]
  217. assert len(parent_states) < 2, "Only one parent state is allowed."
  218. return parent_states[0] if len(parent_states) == 1 else None # type: ignore
  219. @classmethod
  220. @functools.lru_cache()
  221. def get_substates(cls) -> Set[Type[State]]:
  222. """Get the substates of the state.
  223. Returns:
  224. The substates of the state.
  225. """
  226. return set(cls.__subclasses__())
  227. @classmethod
  228. @functools.lru_cache()
  229. def get_name(cls) -> str:
  230. """Get the name of the state.
  231. Returns:
  232. The name of the state.
  233. """
  234. return format.to_snake_case(cls.__name__)
  235. @classmethod
  236. @functools.lru_cache()
  237. def get_full_name(cls) -> str:
  238. """Get the full name of the state.
  239. Returns:
  240. The full name of the state.
  241. """
  242. name = cls.get_name()
  243. parent_state = cls.get_parent_state()
  244. if parent_state is not None:
  245. name = ".".join((parent_state.get_full_name(), name))
  246. return name
  247. @classmethod
  248. @functools.lru_cache()
  249. def get_class_substate(cls, path: Sequence[str]) -> Type[State]:
  250. """Get the class substate.
  251. Args:
  252. path: The path to the substate.
  253. Returns:
  254. The class substate.
  255. Raises:
  256. ValueError: If the substate is not found.
  257. """
  258. if len(path) == 0:
  259. return cls
  260. if path[0] == cls.get_name():
  261. if len(path) == 1:
  262. return cls
  263. path = path[1:]
  264. for substate in cls.get_substates():
  265. if path[0] == substate.get_name():
  266. return substate.get_class_substate(path[1:])
  267. raise ValueError(f"Invalid path: {path}")
  268. @classmethod
  269. def get_class_var(cls, path: Sequence[str]) -> Any:
  270. """Get the class var.
  271. Args:
  272. path: The path to the var.
  273. Returns:
  274. The class var.
  275. Raises:
  276. ValueError: If the path is invalid.
  277. """
  278. path, name = path[:-1], path[-1]
  279. substate = cls.get_class_substate(tuple(path))
  280. if not hasattr(substate, name):
  281. raise ValueError(f"Invalid path: {path}")
  282. return getattr(substate, name)
  283. @classmethod
  284. def _init_var(cls, prop: BaseVar):
  285. """Initialize a variable.
  286. Args:
  287. prop (BaseVar): The variable to initialize
  288. Raises:
  289. TypeError: if the variable has an incorrect type
  290. """
  291. if not types.is_valid_var_type(prop.type_):
  292. raise TypeError(
  293. "State vars must be primitive Python types, "
  294. "Plotly figures, Pandas dataframes, "
  295. "or subclasses of pc.Base. "
  296. f'Found var "{prop.name}" with type {prop.type_}.'
  297. )
  298. cls._set_var(prop)
  299. cls._create_setter(prop)
  300. cls._set_default_value(prop)
  301. @classmethod
  302. def add_var(cls, name: str, type_: Any, default_value: Any = None):
  303. """Add dynamically a variable to the State.
  304. The variable added this way can be used in the same way as a variable
  305. defined statically in the model.
  306. Args:
  307. name: The name of the variable
  308. type_: The type of the variable
  309. default_value: The default value of the variable
  310. Raises:
  311. NameError: if a variable of this name already exists
  312. """
  313. if name in cls.__fields__:
  314. raise NameError(
  315. f"The variable '{name}' already exist. Use a different name"
  316. )
  317. # create the variable based on name and type
  318. var = BaseVar(name=name, type_=type_)
  319. var.set_state(cls)
  320. # add the pydantic field dynamically (must be done before _init_var)
  321. cls.add_field(var, default_value)
  322. cls._init_var(var)
  323. # update the internal dicts so the new variable is correctly handled
  324. cls.base_vars.update({name: var})
  325. cls.vars.update({name: var})
  326. # let substates know about the new variable
  327. for substate_class in cls.__subclasses__():
  328. substate_class.vars.setdefault(name, var)
  329. @classmethod
  330. def _set_var(cls, prop: BaseVar):
  331. """Set the var as a class member.
  332. Args:
  333. prop: The var instance to set.
  334. """
  335. setattr(cls, prop.name, prop)
  336. @classmethod
  337. def _create_setter(cls, prop: BaseVar):
  338. """Create a setter for the var.
  339. Args:
  340. prop: The var to create a setter for.
  341. """
  342. setter_name = prop.get_setter_name(include_state=False)
  343. if setter_name not in cls.__dict__:
  344. event_handler = EventHandler(fn=prop.get_setter())
  345. cls.event_handlers[setter_name] = event_handler
  346. setattr(cls, setter_name, event_handler)
  347. @classmethod
  348. def _set_default_value(cls, prop: BaseVar):
  349. """Set the default value for the var.
  350. Args:
  351. prop: The var to set the default value for.
  352. """
  353. # Get the pydantic field for the var.
  354. field = cls.get_fields()[prop.name]
  355. default_value = prop.get_default_value()
  356. if field.required and default_value is not None:
  357. field.required = False
  358. field.default = default_value
  359. def get_token(self) -> str:
  360. """Return the token of the client associated with this state.
  361. Returns:
  362. The token of the client.
  363. """
  364. return self.router_data.get(constants.RouteVar.CLIENT_TOKEN, "")
  365. def get_sid(self) -> str:
  366. """Return the session ID of the client associated with this state.
  367. Returns:
  368. The session ID of the client.
  369. """
  370. return self.router_data.get(constants.RouteVar.SESSION_ID, "")
  371. def get_headers(self) -> Dict:
  372. """Return the headers of the client associated with this state.
  373. Returns:
  374. The headers of the client.
  375. """
  376. return self.router_data.get(constants.RouteVar.HEADERS, {})
  377. def get_client_ip(self) -> str:
  378. """Return the IP of the client associated with this state.
  379. Returns:
  380. The IP of the client.
  381. """
  382. return self.router_data.get(constants.RouteVar.CLIENT_IP, "")
  383. def get_current_page(self) -> str:
  384. """Obtain the path of current page from the router data.
  385. Returns:
  386. The current page.
  387. """
  388. return self.router_data.get(constants.RouteVar.PATH, "")
  389. def get_query_params(self) -> Dict[str, str]:
  390. """Obtain the query parameters for the queried page.
  391. The query object contains both the URI parameters and the GET parameters.
  392. Returns:
  393. The dict of query parameters.
  394. """
  395. return self.router_data.get(constants.RouteVar.QUERY, {})
  396. def get_cookies(self) -> Dict[str, str]:
  397. """Obtain the cookies of the client stored in the browser.
  398. Returns:
  399. The dict of cookies.
  400. """
  401. headers = self.get_headers().get(constants.RouteVar.COOKIE)
  402. return (
  403. {
  404. pair[0].strip(): pair[1].strip()
  405. for pair in (item.split("=") for item in headers.split(";"))
  406. }
  407. if headers
  408. else {}
  409. )
  410. @classmethod
  411. def setup_dynamic_args(cls, args: dict[str, str]):
  412. """Set up args for easy access in renderer.
  413. Args:
  414. args: a dict of args
  415. """
  416. def argsingle_factory(param):
  417. @ComputedVar
  418. def inner_func(self) -> str:
  419. return self.get_query_params().get(param, "")
  420. return inner_func
  421. def arglist_factory(param):
  422. @ComputedVar
  423. def inner_func(self) -> List:
  424. return self.get_query_params().get(param, [])
  425. return inner_func
  426. for param, value in args.items():
  427. if value == constants.RouteArgType.SINGLE:
  428. func = argsingle_factory(param)
  429. elif value == constants.RouteArgType.LIST:
  430. func = arglist_factory(param)
  431. else:
  432. continue
  433. func.fget.__name__ = param # to allow passing as a prop
  434. cls.vars[param] = cls.computed_vars[param] = func.set_state(cls) # type: ignore
  435. setattr(cls, param, func)
  436. def __getattribute__(self, name: str) -> Any:
  437. """Get the state var.
  438. If the var is inherited, get the var from the parent state.
  439. Args:
  440. name: The name of the var.
  441. Returns:
  442. The value of the var.
  443. """
  444. # If the state hasn't been initialized yet, return the default value.
  445. if not super().__getattribute__("__dict__"):
  446. return super().__getattribute__(name)
  447. inherited_vars = {
  448. **super().__getattribute__("inherited_vars"),
  449. **super().__getattribute__("inherited_backend_vars"),
  450. }
  451. if name in inherited_vars:
  452. return getattr(super().__getattribute__("parent_state"), name)
  453. elif name in super().__getattribute__("_backend_vars"):
  454. return super().__getattribute__("_backend_vars").__getitem__(name)
  455. return super().__getattribute__(name)
  456. def __setattr__(self, name: str, value: Any):
  457. """Set the attribute.
  458. If the attribute is inherited, set the attribute on the parent state.
  459. Args:
  460. name: The name of the attribute.
  461. value: The value of the attribute.
  462. """
  463. # Set the var on the parent state.
  464. inherited_vars = {**self.inherited_vars, **self.inherited_backend_vars}
  465. if name in inherited_vars:
  466. setattr(self.parent_state, name, value)
  467. return
  468. if types.is_backend_variable(name) and name != "_backend_vars":
  469. self._backend_vars.__setitem__(name, value)
  470. self.dirty_vars.add(name)
  471. self.mark_dirty()
  472. return
  473. # Set the attribute.
  474. super().__setattr__(name, value)
  475. # Add the var to the dirty list.
  476. if name in self.vars or name in self.computed_var_dependencies:
  477. self.dirty_vars.add(name)
  478. self.mark_dirty()
  479. # For now, handle router_data updates as a special case
  480. if name == constants.ROUTER_DATA:
  481. self.dirty_vars.add(name)
  482. self.mark_dirty()
  483. # propagate router_data updates down the state tree
  484. for substate in self.substates.values():
  485. setattr(substate, name, value)
  486. def reset(self):
  487. """Reset all the base vars to their default values."""
  488. # Reset the base vars.
  489. fields = self.get_fields()
  490. for prop_name in self.base_vars:
  491. setattr(self, prop_name, fields[prop_name].default)
  492. # Recursively reset the substates.
  493. for substate in self.substates.values():
  494. substate.reset()
  495. # Clean the state.
  496. self.clean()
  497. def get_substate(self, path: Sequence[str]) -> Optional[State]:
  498. """Get the substate.
  499. Args:
  500. path: The path to the substate.
  501. Returns:
  502. The substate.
  503. Raises:
  504. ValueError: If the substate is not found.
  505. """
  506. if len(path) == 0:
  507. return self
  508. if path[0] == self.get_name():
  509. if len(path) == 1:
  510. return self
  511. path = path[1:]
  512. if path[0] not in self.substates:
  513. raise ValueError(f"Invalid path: {path}")
  514. return self.substates[path[0]].get_substate(path[1:])
  515. async def _process(self, event: Event) -> AsyncIterator[StateUpdate]:
  516. """Obtain event info and process event.
  517. Args:
  518. event: The event to process.
  519. Yields:
  520. The state update after processing the event.
  521. Raises:
  522. ValueError: If the state value is None.
  523. """
  524. # Get the event handler.
  525. path = event.name.split(".")
  526. path, name = path[:-1], path[-1]
  527. substate = self.get_substate(path)
  528. handler = substate.event_handlers[name] # type: ignore
  529. if not substate:
  530. raise ValueError(
  531. "The value of state cannot be None when processing an event."
  532. )
  533. # Get the event generator.
  534. event_iter = self._process_event(
  535. handler=handler,
  536. state=substate,
  537. payload=event.payload,
  538. )
  539. # Clean the state before processing the event.
  540. self.clean()
  541. # Run the event generator and return state updates.
  542. async for events, final in event_iter:
  543. # Fix the returned events.
  544. events = fix_events(events, event.token) # type: ignore
  545. # Get the delta after processing the event.
  546. delta = self.get_delta()
  547. # Yield the state update.
  548. yield StateUpdate(delta=delta, events=events, final=final)
  549. # Clean the state to prepare for the next event.
  550. self.clean()
  551. async def _process_event(
  552. self, handler: EventHandler, state: State, payload: Dict
  553. ) -> AsyncIterator[Tuple[Optional[List[EventSpec]], bool]]:
  554. """Process event.
  555. Args:
  556. handler: Eventhandler to process.
  557. state: State to process the handler.
  558. payload: The event payload.
  559. Yields:
  560. Tuple containing:
  561. 0: The state update after processing the event.
  562. 1: Whether the event is the final event.
  563. """
  564. # Get the function to process the event.
  565. fn = functools.partial(handler.fn, state)
  566. # Wrap the function in a try/except block.
  567. try:
  568. # Handle async functions.
  569. if asyncio.iscoroutinefunction(fn.func):
  570. events = await fn(**payload)
  571. # Handle regular functions.
  572. else:
  573. events = fn(**payload)
  574. # Handle async generators.
  575. if inspect.isasyncgen(events):
  576. async for event in events:
  577. yield event, False
  578. yield None, True
  579. # Handle regular generators.
  580. elif inspect.isgenerator(events):
  581. for event in events:
  582. yield event, False
  583. yield None, True
  584. # Handle regular event chains.
  585. else:
  586. yield events, True
  587. # If an error occurs, throw a window alert.
  588. except Exception:
  589. error = traceback.format_exc()
  590. print(error)
  591. yield [window_alert("An error occurred. See logs for details.")], True
  592. def _always_dirty_computed_vars(self) -> Set[str]:
  593. """The set of ComputedVars that always need to be recalculated.
  594. Returns:
  595. Set of all ComputedVar in this state where cache=False
  596. """
  597. return set(
  598. cvar_name
  599. for cvar_name, cvar in self.computed_vars.items()
  600. if not cvar.cache
  601. )
  602. def _mark_dirty_computed_vars(self) -> None:
  603. """Mark ComputedVars that need to be recalculated based on dirty_vars."""
  604. dirty_vars = self.dirty_vars
  605. while dirty_vars:
  606. calc_vars, dirty_vars = dirty_vars, set()
  607. for cvar in self._dirty_computed_vars(from_vars=calc_vars):
  608. self.dirty_vars.add(cvar)
  609. dirty_vars.add(cvar)
  610. actual_var = self.computed_vars.get(cvar)
  611. if actual_var:
  612. actual_var.mark_dirty(instance=self)
  613. def _dirty_computed_vars(self, from_vars: Optional[Set[str]] = None) -> Set[str]:
  614. """Determine ComputedVars that need to be recalculated based on the given vars.
  615. Args:
  616. from_vars: find ComputedVar that depend on this set of vars. If unspecified, will use the dirty_vars.
  617. Returns:
  618. Set of computed vars to include in the delta.
  619. """
  620. return set(
  621. cvar
  622. for dirty_var in from_vars or self.dirty_vars
  623. for cvar in self.computed_var_dependencies[dirty_var]
  624. )
  625. def get_delta(self) -> Delta:
  626. """Get the delta for the state.
  627. Returns:
  628. The delta for the state.
  629. """
  630. delta = {}
  631. # Apply dirty variables down into substates
  632. self.dirty_vars.update(self._always_dirty_computed_vars())
  633. self.mark_dirty()
  634. # Return the dirty vars for this instance, any cached/dependent computed vars,
  635. # and always dirty computed vars (cache=False)
  636. delta_vars = (
  637. self.dirty_vars.intersection(self.base_vars)
  638. .union(self._dirty_computed_vars())
  639. .union(self._always_dirty_computed_vars())
  640. )
  641. subdelta = {
  642. prop: getattr(self, prop)
  643. for prop in delta_vars
  644. if not types.is_backend_variable(prop)
  645. }
  646. if len(subdelta) > 0:
  647. delta[self.get_full_name()] = subdelta
  648. # Recursively find the substate deltas.
  649. substates = self.substates
  650. for substate in self.dirty_substates:
  651. delta.update(substates[substate].get_delta())
  652. # Format the delta.
  653. delta = format.format_state(delta)
  654. # Return the delta.
  655. return delta
  656. def mark_dirty(self):
  657. """Mark the substate and all parent states as dirty."""
  658. state_name = self.get_name()
  659. if (
  660. self.parent_state is not None
  661. and state_name not in self.parent_state.dirty_substates
  662. ):
  663. self.parent_state.dirty_substates.add(self.get_name())
  664. self.parent_state.mark_dirty()
  665. # have to mark computed vars dirty to allow access to newly computed
  666. # values within the same ComputedVar function
  667. self._mark_dirty_computed_vars()
  668. # Propagate dirty var / computed var status into substates
  669. substates = self.substates
  670. for var in self.dirty_vars:
  671. for substate_name in self.substate_var_dependencies[var]:
  672. self.dirty_substates.add(substate_name)
  673. substate = substates[substate_name]
  674. substate.dirty_vars.add(var)
  675. substate.mark_dirty()
  676. def clean(self):
  677. """Reset the dirty vars."""
  678. # Recursively clean the substates.
  679. for substate in self.dirty_substates:
  680. self.substates[substate].clean()
  681. # Clean this state.
  682. self.dirty_vars = set()
  683. self.dirty_substates = set()
  684. def dict(self, include_computed: bool = True, **kwargs) -> Dict[str, Any]:
  685. """Convert the object to a dictionary.
  686. Args:
  687. include_computed: Whether to include computed vars.
  688. **kwargs: Kwargs to pass to the pydantic dict method.
  689. Returns:
  690. The object as a dictionary.
  691. """
  692. if include_computed:
  693. # Apply dirty variables down into substates to allow never-cached ComputedVar to
  694. # trigger recalculation of dependent vars
  695. self.dirty_vars.update(self._always_dirty_computed_vars())
  696. self.mark_dirty()
  697. base_vars = {
  698. prop_name: self.get_value(getattr(self, prop_name))
  699. for prop_name in self.base_vars
  700. }
  701. computed_vars = (
  702. {
  703. # Include the computed vars.
  704. prop_name: self.get_value(getattr(self, prop_name))
  705. for prop_name in self.computed_vars
  706. }
  707. if include_computed
  708. else {}
  709. )
  710. substate_vars = {
  711. k: v.dict(include_computed=include_computed, **kwargs)
  712. for k, v in self.substates.items()
  713. }
  714. variables = {**base_vars, **computed_vars, **substate_vars}
  715. return {k: variables[k] for k in sorted(variables)}
  716. class DefaultState(State):
  717. """The default empty state."""
  718. pass
  719. class StateUpdate(Base):
  720. """A state update sent to the frontend."""
  721. # The state delta.
  722. delta: Delta = {}
  723. # Events to be added to the event queue.
  724. events: List[Event] = []
  725. # Whether this is the final state update for the event.
  726. final: bool = True
  727. class StateManager(Base):
  728. """A class to manage many client states."""
  729. # The state class to use.
  730. state: Type[State] = DefaultState
  731. # The mapping of client ids to states.
  732. states: Dict[str, State] = {}
  733. # The token expiration time (s).
  734. token_expiration: int = constants.TOKEN_EXPIRATION
  735. # The redis client to use.
  736. redis: Optional[Redis] = None
  737. def setup(self, state: Type[State]):
  738. """Set up the state manager.
  739. Args:
  740. state: The state class to use.
  741. """
  742. self.state = state
  743. self.redis = prerequisites.get_redis()
  744. def get_state(self, token: str) -> State:
  745. """Get the state for a token.
  746. Args:
  747. token: The token to get the state for.
  748. Returns:
  749. The state for the token.
  750. """
  751. if self.redis is not None:
  752. redis_state = self.redis.get(token)
  753. if redis_state is None:
  754. self.set_state(token, self.state())
  755. return self.get_state(token)
  756. return cloudpickle.loads(redis_state)
  757. if token not in self.states:
  758. self.states[token] = self.state()
  759. return self.states[token]
  760. def set_state(self, token: str, state: State):
  761. """Set the state for a token.
  762. Args:
  763. token: The token to set the state for.
  764. state: The state to set.
  765. """
  766. if self.redis is None:
  767. return
  768. self.redis.set(token, cloudpickle.dumps(state), ex=self.token_expiration)
  769. def _convert_mutable_datatypes(
  770. field_value: Any, reassign_field: Callable, field_name: str
  771. ) -> Any:
  772. """Recursively convert mutable data to the Pc data types.
  773. Note: right now only list & dict would be handled recursively.
  774. Args:
  775. field_value: The target field_value.
  776. reassign_field:
  777. The function to reassign the field in the parent state.
  778. field_name: the name of the field in the parent state
  779. Returns:
  780. The converted field_value
  781. """
  782. if isinstance(field_value, list):
  783. for index in range(len(field_value)):
  784. field_value[index] = _convert_mutable_datatypes(
  785. field_value[index], reassign_field, field_name
  786. )
  787. field_value = PCList(
  788. field_value, reassign_field=reassign_field, field_name=field_name
  789. )
  790. if isinstance(field_value, dict):
  791. for key, value in field_value.items():
  792. field_value[key] = _convert_mutable_datatypes(
  793. value, reassign_field, field_name
  794. )
  795. field_value = PCDict(
  796. field_value, reassign_field=reassign_field, field_name=field_name
  797. )
  798. return field_value