state.py 132 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849285028512852285328542855285628572858285928602861286228632864286528662867286828692870287128722873287428752876287728782879288028812882288328842885288628872888288928902891289228932894289528962897289828992900290129022903290429052906290729082909291029112912291329142915291629172918291929202921292229232924292529262927292829292930293129322933293429352936293729382939294029412942294329442945294629472948294929502951295229532954295529562957295829592960296129622963296429652966296729682969297029712972297329742975297629772978297929802981298229832984298529862987298829892990299129922993299429952996299729982999300030013002300330043005300630073008300930103011301230133014301530163017301830193020302130223023302430253026302730283029303030313032303330343035303630373038303930403041304230433044304530463047304830493050305130523053305430553056305730583059306030613062306330643065306630673068306930703071307230733074307530763077307830793080308130823083308430853086308730883089309030913092309330943095309630973098309931003101310231033104310531063107310831093110311131123113311431153116311731183119312031213122312331243125312631273128312931303131313231333134313531363137313831393140314131423143314431453146314731483149315031513152315331543155315631573158315931603161316231633164316531663167316831693170317131723173317431753176317731783179318031813182318331843185318631873188318931903191319231933194319531963197319831993200320132023203320432053206320732083209321032113212321332143215321632173218321932203221322232233224322532263227322832293230323132323233323432353236323732383239324032413242324332443245324632473248324932503251325232533254325532563257325832593260326132623263326432653266326732683269327032713272327332743275327632773278327932803281328232833284328532863287328832893290329132923293329432953296329732983299330033013302330333043305330633073308330933103311331233133314331533163317331833193320332133223323332433253326332733283329333033313332333333343335333633373338333933403341334233433344334533463347334833493350335133523353335433553356335733583359336033613362336333643365336633673368336933703371337233733374337533763377337833793380338133823383338433853386338733883389339033913392339333943395339633973398339934003401340234033404340534063407340834093410341134123413341434153416341734183419342034213422342334243425342634273428342934303431343234333434343534363437343834393440344134423443344434453446344734483449345034513452345334543455345634573458345934603461346234633464346534663467346834693470347134723473347434753476347734783479348034813482348334843485348634873488348934903491349234933494349534963497349834993500350135023503350435053506350735083509351035113512351335143515351635173518351935203521352235233524352535263527352835293530353135323533353435353536353735383539354035413542354335443545354635473548354935503551355235533554355535563557355835593560356135623563356435653566356735683569357035713572357335743575357635773578357935803581358235833584358535863587358835893590359135923593359435953596359735983599360036013602360336043605360636073608360936103611361236133614361536163617361836193620362136223623362436253626362736283629363036313632363336343635363636373638363936403641364236433644364536463647364836493650365136523653365436553656365736583659366036613662366336643665366636673668366936703671367236733674367536763677367836793680368136823683368436853686368736883689369036913692369336943695369636973698369937003701370237033704370537063707370837093710371137123713371437153716371737183719372037213722372337243725372637273728372937303731373237333734373537363737373837393740374137423743374437453746374737483749375037513752375337543755375637573758375937603761376237633764376537663767376837693770377137723773377437753776377737783779378037813782378337843785378637873788378937903791379237933794379537963797379837993800380138023803380438053806380738083809381038113812381338143815381638173818381938203821382238233824382538263827382838293830383138323833383438353836383738383839384038413842384338443845384638473848384938503851385238533854385538563857385838593860386138623863386438653866386738683869387038713872387338743875387638773878387938803881388238833884388538863887388838893890389138923893389438953896389738983899390039013902390339043905390639073908390939103911391239133914391539163917391839193920392139223923392439253926392739283929393039313932393339343935393639373938393939403941394239433944394539463947394839493950395139523953395439553956395739583959396039613962396339643965396639673968396939703971397239733974
  1. """Define the reflex state specification."""
  2. from __future__ import annotations
  3. import asyncio
  4. import contextlib
  5. import copy
  6. import dataclasses
  7. import functools
  8. import inspect
  9. import json
  10. import pickle
  11. import sys
  12. import typing
  13. import uuid
  14. from abc import ABC, abstractmethod
  15. from collections import defaultdict
  16. from hashlib import md5
  17. from pathlib import Path
  18. from types import FunctionType, MethodType
  19. from typing import (
  20. TYPE_CHECKING,
  21. Any,
  22. AsyncIterator,
  23. BinaryIO,
  24. Callable,
  25. ClassVar,
  26. Dict,
  27. List,
  28. Optional,
  29. Sequence,
  30. Set,
  31. Tuple,
  32. Type,
  33. TypeVar,
  34. Union,
  35. cast,
  36. get_args,
  37. get_type_hints,
  38. )
  39. from sqlalchemy.orm import DeclarativeBase
  40. from typing_extensions import Self
  41. from reflex import event
  42. from reflex.config import PerformanceMode, get_config
  43. from reflex.istate.data import RouterData
  44. from reflex.istate.storage import ClientStorageBase
  45. from reflex.model import Model
  46. from reflex.vars.base import (
  47. ComputedVar,
  48. DynamicRouteVar,
  49. Var,
  50. computed_var,
  51. dispatch,
  52. get_unique_variable_name,
  53. is_computed_var,
  54. )
  55. try:
  56. import pydantic.v1 as pydantic
  57. except ModuleNotFoundError:
  58. import pydantic
  59. from pydantic import BaseModel as BaseModelV2
  60. try:
  61. from pydantic.v1 import BaseModel as BaseModelV1
  62. except ModuleNotFoundError:
  63. BaseModelV1 = BaseModelV2
  64. import wrapt
  65. from redis.asyncio import Redis
  66. from redis.exceptions import ResponseError
  67. import reflex.istate.dynamic
  68. from reflex import constants
  69. from reflex.base import Base
  70. from reflex.config import environment
  71. from reflex.event import (
  72. BACKGROUND_TASK_MARKER,
  73. Event,
  74. EventHandler,
  75. EventSpec,
  76. fix_events,
  77. )
  78. from reflex.utils import console, format, path_ops, prerequisites, types
  79. from reflex.utils.exceptions import (
  80. ComputedVarShadowsBaseVars,
  81. ComputedVarShadowsStateVar,
  82. DynamicComponentInvalidSignature,
  83. DynamicRouteArgShadowsStateVar,
  84. EventHandlerShadowsBuiltInStateMethod,
  85. ImmutableStateError,
  86. InvalidStateManagerMode,
  87. LockExpiredError,
  88. ReflexRuntimeError,
  89. SetUndefinedStateVarError,
  90. StateSchemaMismatchError,
  91. StateTooLargeError,
  92. )
  93. from reflex.utils.exec import is_testing_env
  94. from reflex.utils.serializers import serializer
  95. from reflex.utils.types import (
  96. _isinstance,
  97. get_origin,
  98. is_union,
  99. override,
  100. value_inside_optional,
  101. )
  102. from reflex.vars import VarData
  103. if TYPE_CHECKING:
  104. from reflex.components.component import Component
  105. Delta = Dict[str, Any]
  106. var = computed_var
  107. if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
  108. # If the state is this large, it's considered a performance issue.
  109. TOO_LARGE_SERIALIZED_STATE = environment.REFLEX_STATE_SIZE_LIMIT.get() * 1024
  110. # Only warn about each state class size once.
  111. _WARNED_ABOUT_STATE_SIZE: Set[str] = set()
  112. # Errors caught during pickling of state
  113. HANDLED_PICKLE_ERRORS = (
  114. pickle.PicklingError,
  115. AttributeError,
  116. IndexError,
  117. TypeError,
  118. ValueError,
  119. )
  120. def _no_chain_background_task(
  121. state_cls: Type["BaseState"], name: str, fn: Callable
  122. ) -> Callable:
  123. """Protect against directly chaining a background task from another event handler.
  124. Args:
  125. state_cls: The state class that the event handler is in.
  126. name: The name of the background task.
  127. fn: The background task coroutine function / generator.
  128. Returns:
  129. A compatible coroutine function / generator that raises a runtime error.
  130. Raises:
  131. TypeError: If the background task is not async.
  132. """
  133. call = f"{state_cls.__name__}.{name}"
  134. message = (
  135. f"Cannot directly call background task {name!r}, use "
  136. f"`yield {call}` or `return {call}` instead."
  137. )
  138. if inspect.iscoroutinefunction(fn):
  139. async def _no_chain_background_task_co(*args, **kwargs):
  140. raise RuntimeError(message)
  141. return _no_chain_background_task_co
  142. if inspect.isasyncgenfunction(fn):
  143. async def _no_chain_background_task_gen(*args, **kwargs):
  144. yield
  145. raise RuntimeError(message)
  146. return _no_chain_background_task_gen
  147. raise TypeError(f"{fn} is marked as a background task, but is not async.")
  148. def _substate_key(
  149. token: str,
  150. state_cls_or_name: BaseState | Type[BaseState] | str | Sequence[str],
  151. ) -> str:
  152. """Get the substate key.
  153. Args:
  154. token: The token of the state.
  155. state_cls_or_name: The state class/instance or name or sequence of name parts.
  156. Returns:
  157. The substate key.
  158. """
  159. if isinstance(state_cls_or_name, BaseState) or (
  160. isinstance(state_cls_or_name, type) and issubclass(state_cls_or_name, BaseState)
  161. ):
  162. state_cls_or_name = state_cls_or_name.get_full_name()
  163. elif isinstance(state_cls_or_name, (list, tuple)):
  164. state_cls_or_name = ".".join(state_cls_or_name)
  165. return f"{token}_{state_cls_or_name}"
  166. def _split_substate_key(substate_key: str) -> tuple[str, str]:
  167. """Split the substate key into token and state name.
  168. Args:
  169. substate_key: The substate key.
  170. Returns:
  171. Tuple of token and state name.
  172. """
  173. token, _, state_name = substate_key.partition("_")
  174. return token, state_name
  175. @dataclasses.dataclass(frozen=True, init=False)
  176. class EventHandlerSetVar(EventHandler):
  177. """A special event handler to wrap setvar functionality."""
  178. state_cls: Type[BaseState] = dataclasses.field(init=False)
  179. def __init__(self, state_cls: Type[BaseState]):
  180. """Initialize the EventHandlerSetVar.
  181. Args:
  182. state_cls: The state class that vars will be set on.
  183. """
  184. super().__init__(
  185. fn=type(self).setvar,
  186. state_full_name=state_cls.get_full_name(),
  187. )
  188. object.__setattr__(self, "state_cls", state_cls)
  189. def setvar(self, var_name: str, value: Any):
  190. """Set the state variable to the value of the event.
  191. Note: `self` here will be an instance of the state, not EventHandlerSetVar.
  192. Args:
  193. var_name: The name of the variable to set.
  194. value: The value to set the variable to.
  195. """
  196. getattr(self, constants.SETTER_PREFIX + var_name)(value)
  197. def __call__(self, *args: Any) -> EventSpec:
  198. """Performs pre-checks and munging on the provided args that will become an EventSpec.
  199. Args:
  200. *args: The event args.
  201. Returns:
  202. The (partial) EventSpec that will be used to create the event to setvar.
  203. Raises:
  204. AttributeError: If the given Var name does not exist on the state.
  205. EventHandlerValueError: If the given Var name is not a str
  206. NotImplementedError: If the setter for the given Var is async
  207. """
  208. from reflex.utils.exceptions import EventHandlerValueError
  209. if args:
  210. if not isinstance(args[0], str):
  211. raise EventHandlerValueError(
  212. f"Var name must be passed as a string, got {args[0]!r}"
  213. )
  214. handler = getattr(self.state_cls, constants.SETTER_PREFIX + args[0], None)
  215. # Check that the requested Var setter exists on the State at compile time.
  216. if handler is None:
  217. raise AttributeError(
  218. f"Variable `{args[0]}` cannot be set on `{self.state_cls.get_full_name()}`"
  219. )
  220. if asyncio.iscoroutinefunction(handler.fn):
  221. raise NotImplementedError(
  222. f"Setter for {args[0]} is async, which is not supported."
  223. )
  224. return super().__call__(*args)
  225. if TYPE_CHECKING:
  226. from pydantic.v1.fields import ModelField
  227. def get_var_for_field(cls: Type[BaseState], f: ModelField):
  228. """Get a Var instance for a Pydantic field.
  229. Args:
  230. cls: The state class.
  231. f: The Pydantic field.
  232. Returns:
  233. The Var instance.
  234. """
  235. from reflex.vars import Field
  236. field_name = format.format_state_name(cls.get_full_name()) + "." + f.name
  237. return dispatch(
  238. field_name=field_name,
  239. var_data=VarData.from_state(cls, f.name),
  240. result_var_type=f.outer_type_
  241. if get_origin(f.outer_type_) is not Field
  242. else get_args(f.outer_type_)[0],
  243. )
  244. # Keep track of all state instances to calculate minified state names
  245. state_count: int = 0
  246. minified_state_names: Dict[str, str] = {}
  247. def next_minified_state_name() -> str:
  248. """Get the next minified state name.
  249. Returns:
  250. The next minified state name.
  251. """
  252. global state_count
  253. num = state_count
  254. # All possible chars for minified state name
  255. chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ$_"
  256. base = len(chars)
  257. state_name = ""
  258. if num == 0:
  259. state_name = chars[0]
  260. while num > 0:
  261. state_name = chars[num % base] + state_name
  262. num = num // base
  263. state_count += 1
  264. return state_name
  265. def get_minified_state_name(state_name: str) -> str:
  266. """Generate a minified state name.
  267. Args:
  268. state_name: The state name to minify.
  269. Returns:
  270. The minified state name.
  271. Raises:
  272. ValueError: If no more minified state names are available
  273. """
  274. if state_name in minified_state_names:
  275. return minified_state_names[state_name]
  276. while name := next_minified_state_name():
  277. if name in minified_state_names.values():
  278. continue
  279. minified_state_names[state_name] = name
  280. return name
  281. raise ValueError("No more minified state names available")
  282. class BaseState(Base, ABC, extra=pydantic.Extra.allow):
  283. """The state of the app."""
  284. # A map from the var name to the var.
  285. vars: ClassVar[Dict[str, Var]] = {}
  286. # The base vars of the class.
  287. base_vars: ClassVar[Dict[str, Var]] = {}
  288. # The computed vars of the class.
  289. computed_vars: ClassVar[Dict[str, ComputedVar]] = {}
  290. # Vars inherited by the parent state.
  291. inherited_vars: ClassVar[Dict[str, Var]] = {}
  292. # Backend base vars that are never sent to the client.
  293. backend_vars: ClassVar[Dict[str, Any]] = {}
  294. # Backend base vars inherited
  295. inherited_backend_vars: ClassVar[Dict[str, Any]] = {}
  296. # The event handlers.
  297. event_handlers: ClassVar[Dict[str, EventHandler]] = {}
  298. # A set of subclassses of this class.
  299. class_subclasses: ClassVar[Set[Type[BaseState]]] = set()
  300. # Mapping of var name to set of computed variables that depend on it
  301. _computed_var_dependencies: ClassVar[Dict[str, Set[str]]] = {}
  302. # Mapping of var name to set of substates that depend on it
  303. _substate_var_dependencies: ClassVar[Dict[str, Set[str]]] = {}
  304. # Set of vars which always need to be recomputed
  305. _always_dirty_computed_vars: ClassVar[Set[str]] = set()
  306. # Set of substates which always need to be recomputed
  307. _always_dirty_substates: ClassVar[Set[str]] = set()
  308. # The parent state.
  309. parent_state: Optional[BaseState] = None
  310. # The substates of the state.
  311. substates: Dict[str, BaseState] = {}
  312. # The set of dirty vars.
  313. dirty_vars: Set[str] = set()
  314. # The set of dirty substates.
  315. dirty_substates: Set[str] = set()
  316. # The routing path that triggered the state
  317. router_data: Dict[str, Any] = {}
  318. # Per-instance copy of backend base variable values
  319. _backend_vars: Dict[str, Any] = {}
  320. # The router data for the current page
  321. router: RouterData = RouterData()
  322. # Whether the state has ever been touched since instantiation.
  323. _was_touched: bool = False
  324. # Whether this state class is a mixin and should not be instantiated.
  325. _mixin: ClassVar[bool] = False
  326. # A special event handler for setting base vars.
  327. setvar: ClassVar[EventHandler]
  328. def __init__(
  329. self,
  330. parent_state: BaseState | None = None,
  331. init_substates: bool = True,
  332. _reflex_internal_init: bool = False,
  333. **kwargs,
  334. ):
  335. """Initialize the state.
  336. DO NOT INSTANTIATE STATE CLASSES DIRECTLY! Use StateManager.get_state() instead.
  337. Args:
  338. parent_state: The parent state.
  339. init_substates: Whether to initialize the substates in this instance.
  340. _reflex_internal_init: A flag to indicate that the state is being initialized by the framework.
  341. **kwargs: The kwargs to set as attributes on the state.
  342. Raises:
  343. ReflexRuntimeError: If the state is instantiated directly by end user.
  344. """
  345. from reflex.utils.exceptions import ReflexRuntimeError
  346. if not _reflex_internal_init and not is_testing_env():
  347. raise ReflexRuntimeError(
  348. "State classes should not be instantiated directly in a Reflex app. "
  349. "See https://reflex.dev/docs/state/ for further information."
  350. )
  351. if type(self)._mixin:
  352. raise ReflexRuntimeError(
  353. f"{type(self).__name__} is a state mixin and cannot be instantiated directly."
  354. )
  355. kwargs["parent_state"] = parent_state
  356. super().__init__()
  357. for name, value in kwargs.items():
  358. setattr(self, name, value)
  359. # Setup the substates (for memory state manager only).
  360. if init_substates:
  361. for substate in self.get_substates():
  362. self.substates[substate.get_name()] = substate(
  363. parent_state=self,
  364. _reflex_internal_init=True,
  365. )
  366. # Create a fresh copy of the backend variables for this instance
  367. self._backend_vars = copy.deepcopy(
  368. {name: item for name, item in self.backend_vars.items()}
  369. )
  370. def __repr__(self) -> str:
  371. """Get the string representation of the state.
  372. Returns:
  373. The string representation of the state.
  374. """
  375. return f"{self.__class__.__name__}({self.dict()})"
  376. @classmethod
  377. def _get_computed_vars(cls) -> list[ComputedVar]:
  378. """Helper function to get all computed vars of a instance.
  379. Returns:
  380. A list of computed vars.
  381. """
  382. return [
  383. v
  384. for mixin in cls._mixins() + [cls]
  385. for name, v in mixin.__dict__.items()
  386. if is_computed_var(v) and name not in cls.inherited_vars
  387. ]
  388. @classmethod
  389. def _validate_module_name(cls) -> None:
  390. """Check if the module name is valid.
  391. Reflex uses ___ as state name module separator.
  392. Raises:
  393. NameError: If the module name is invalid.
  394. """
  395. if "___" in cls.__module__:
  396. raise NameError(
  397. "The module name of a State class cannot contain '___'. "
  398. "Please rename the module."
  399. )
  400. @classmethod
  401. def __init_subclass__(cls, mixin: bool = False, **kwargs):
  402. """Do some magic for the subclass initialization.
  403. Args:
  404. mixin: Whether the subclass is a mixin and should not be initialized.
  405. **kwargs: The kwargs to pass to the pydantic init_subclass method.
  406. Raises:
  407. StateValueError: If a substate class shadows another.
  408. """
  409. from reflex.utils.exceptions import StateValueError
  410. super().__init_subclass__(**kwargs)
  411. cls._mixin = mixin
  412. if mixin:
  413. return
  414. # Handle locally-defined states for pickling.
  415. if "<locals>" in cls.__qualname__:
  416. cls._handle_local_def()
  417. # Validate the module name.
  418. cls._validate_module_name()
  419. # Event handlers should not shadow builtin state methods.
  420. cls._check_overridden_methods()
  421. # Computed vars should not shadow builtin state props.
  422. cls._check_overridden_basevars()
  423. # Reset subclass tracking for this class.
  424. cls.class_subclasses = set()
  425. # Reset dirty substate tracking for this class.
  426. cls._always_dirty_substates = set()
  427. # Get the parent vars.
  428. parent_state = cls.get_parent_state()
  429. if parent_state is not None:
  430. cls.inherited_vars = parent_state.vars
  431. cls.inherited_backend_vars = parent_state.backend_vars
  432. # Check if another substate class with the same name has already been defined.
  433. if cls.get_name() in set(
  434. c.get_name() for c in parent_state.class_subclasses
  435. ):
  436. # This should not happen, since we have added module prefix to state names in #3214
  437. raise StateValueError(
  438. f"The substate class '{cls.get_name()}' has been defined multiple times. "
  439. "Shadowing substate classes is not allowed."
  440. )
  441. # Track this new subclass in the parent state's subclasses set.
  442. parent_state.class_subclasses.add(cls)
  443. # Get computed vars.
  444. computed_vars = cls._get_computed_vars()
  445. cls._check_overridden_computed_vars()
  446. new_backend_vars = {
  447. name: value
  448. for name, value in cls.__dict__.items()
  449. if types.is_backend_base_variable(name, cls)
  450. }
  451. # Add annotated backend vars that may not have a default value.
  452. new_backend_vars.update(
  453. {
  454. name: cls._get_var_default(name, annotation_value)
  455. for name, annotation_value in cls._get_type_hints().items()
  456. if name not in new_backend_vars
  457. and types.is_backend_base_variable(name, cls)
  458. }
  459. )
  460. cls.backend_vars = {
  461. **cls.inherited_backend_vars,
  462. **new_backend_vars,
  463. }
  464. # Set the base and computed vars.
  465. cls.base_vars = {
  466. f.name: get_var_for_field(cls, f)
  467. for f in cls.get_fields().values()
  468. if f.name not in cls.get_skip_vars()
  469. }
  470. cls.computed_vars = {
  471. v._js_expr: v._replace(merge_var_data=VarData.from_state(cls))
  472. for v in computed_vars
  473. }
  474. cls.vars = {
  475. **cls.inherited_vars,
  476. **cls.base_vars,
  477. **cls.computed_vars,
  478. }
  479. cls.event_handlers = {}
  480. # Setup the base vars at the class level.
  481. for prop in cls.base_vars.values():
  482. cls._init_var(prop)
  483. # Set up the event handlers.
  484. events = {
  485. name: fn
  486. for name, fn in cls.__dict__.items()
  487. if cls._item_is_event_handler(name, fn)
  488. }
  489. for mixin in cls._mixins():
  490. for name, value in mixin.__dict__.items():
  491. if name in cls.inherited_vars:
  492. continue
  493. if is_computed_var(value):
  494. fget = cls._copy_fn(value.fget)
  495. newcv = value._replace(fget=fget, _var_data=VarData.from_state(cls))
  496. # cleanup refs to mixin cls in var_data
  497. setattr(cls, name, newcv)
  498. cls.computed_vars[newcv._js_expr] = newcv
  499. cls.vars[newcv._js_expr] = newcv
  500. continue
  501. if types.is_backend_base_variable(name, mixin):
  502. cls.backend_vars[name] = copy.deepcopy(value)
  503. continue
  504. if events.get(name) is not None:
  505. continue
  506. if not cls._item_is_event_handler(name, value):
  507. continue
  508. if parent_state is not None and parent_state.event_handlers.get(name):
  509. continue
  510. value = cls._copy_fn(value)
  511. value.__qualname__ = f"{cls.__name__}.{name}"
  512. events[name] = value
  513. # Create the setvar event handler for this state
  514. cls._create_setvar()
  515. for name, fn in events.items():
  516. handler = cls._create_event_handler(fn)
  517. cls.event_handlers[name] = handler
  518. setattr(cls, name, handler)
  519. # Initialize per-class var dependency tracking.
  520. cls._computed_var_dependencies = defaultdict(set)
  521. cls._substate_var_dependencies = defaultdict(set)
  522. cls._init_var_dependency_dicts()
  523. @staticmethod
  524. def _copy_fn(fn: Callable) -> Callable:
  525. """Copy a function. Used to copy ComputedVars and EventHandlers from mixins.
  526. Args:
  527. fn: The function to copy.
  528. Returns:
  529. The copied function.
  530. """
  531. newfn = FunctionType(
  532. fn.__code__,
  533. fn.__globals__,
  534. name=fn.__name__,
  535. argdefs=fn.__defaults__,
  536. closure=fn.__closure__,
  537. )
  538. newfn.__annotations__ = fn.__annotations__
  539. if mark := getattr(fn, BACKGROUND_TASK_MARKER, None):
  540. setattr(newfn, BACKGROUND_TASK_MARKER, mark)
  541. return newfn
  542. @staticmethod
  543. def _item_is_event_handler(name: str, value: Any) -> bool:
  544. """Check if the item is an event handler.
  545. Args:
  546. name: The name of the item.
  547. value: The value of the item.
  548. Returns:
  549. Whether the item is an event handler.
  550. """
  551. return (
  552. not name.startswith("_")
  553. and isinstance(value, Callable)
  554. and not isinstance(value, EventHandler)
  555. and hasattr(value, "__code__")
  556. )
  557. @classmethod
  558. def _evaluate(
  559. cls, f: Callable[[Self], Any], of_type: Union[type, None] = None
  560. ) -> Var:
  561. """Evaluate a function to a ComputedVar. Experimental.
  562. Args:
  563. f: The function to evaluate.
  564. of_type: The type of the ComputedVar. Defaults to Component.
  565. Returns:
  566. The ComputedVar.
  567. """
  568. console.warn(
  569. "The _evaluate method is experimental and may be removed in future versions."
  570. )
  571. from reflex.components.component import Component
  572. of_type = of_type or Component
  573. unique_var_name = get_unique_variable_name()
  574. @computed_var(_js_expr=unique_var_name, return_type=of_type)
  575. def computed_var_func(state: Self):
  576. result = f(state)
  577. if not _isinstance(result, of_type):
  578. console.warn(
  579. f"Inline ComputedVar {f} expected type {of_type}, got {type(result)}. "
  580. "You can specify expected type with `of_type` argument."
  581. )
  582. return result
  583. setattr(cls, unique_var_name, computed_var_func)
  584. cls.computed_vars[unique_var_name] = computed_var_func
  585. cls.vars[unique_var_name] = computed_var_func
  586. cls._update_substate_inherited_vars({unique_var_name: computed_var_func})
  587. cls._always_dirty_computed_vars.add(unique_var_name)
  588. return getattr(cls, unique_var_name)
  589. @classmethod
  590. def _mixins(cls) -> List[Type]:
  591. """Get the mixin classes of the state.
  592. Returns:
  593. The mixin classes of the state.
  594. """
  595. return [
  596. mixin
  597. for mixin in cls.__mro__
  598. if (
  599. mixin not in [pydantic.BaseModel, Base, cls]
  600. and issubclass(mixin, BaseState)
  601. and mixin._mixin is True
  602. )
  603. ]
  604. @classmethod
  605. def _handle_local_def(cls):
  606. """Handle locally-defined states for pickling."""
  607. known_names = dir(reflex.istate.dynamic)
  608. proposed_name = cls.__name__
  609. for ix in range(len(known_names)):
  610. if proposed_name not in known_names:
  611. break
  612. proposed_name = f"{cls.__name__}_{ix}"
  613. setattr(reflex.istate.dynamic, proposed_name, cls)
  614. cls.__original_name__ = cls.__name__
  615. cls.__original_module__ = cls.__module__
  616. cls.__name__ = cls.__qualname__ = proposed_name
  617. cls.__module__ = reflex.istate.dynamic.__name__
  618. @classmethod
  619. def _get_type_hints(cls) -> dict[str, Any]:
  620. """Get the type hints for this class.
  621. If the class is dynamic, evaluate the type hints with the original
  622. module in the local namespace.
  623. Returns:
  624. The type hints dict.
  625. """
  626. original_module = getattr(cls, "__original_module__", None)
  627. if original_module is not None:
  628. localns = sys.modules[original_module].__dict__
  629. else:
  630. localns = None
  631. return get_type_hints(cls, localns=localns)
  632. @classmethod
  633. def _init_var_dependency_dicts(cls):
  634. """Initialize the var dependency tracking dicts.
  635. Allows the state to know which vars each ComputedVar depends on and
  636. whether a ComputedVar depends on a var in its parent state.
  637. Additional updates tracking dicts for vars and substates that always
  638. need to be recomputed.
  639. """
  640. inherited_vars = set(cls.inherited_vars).union(
  641. set(cls.inherited_backend_vars),
  642. )
  643. for cvar_name, cvar in cls.computed_vars.items():
  644. # Add the dependencies.
  645. for var in cvar._deps(objclass=cls):
  646. cls._computed_var_dependencies[var].add(cvar_name)
  647. if var in inherited_vars:
  648. # track that this substate depends on its parent for this var
  649. state_name = cls.get_name()
  650. parent_state = cls.get_parent_state()
  651. while parent_state is not None and var in {
  652. **parent_state.vars,
  653. **parent_state.backend_vars,
  654. }:
  655. parent_state._substate_var_dependencies[var].add(state_name)
  656. state_name, parent_state = (
  657. parent_state.get_name(),
  658. parent_state.get_parent_state(),
  659. )
  660. # ComputedVar with cache=False always need to be recomputed
  661. cls._always_dirty_computed_vars = set(
  662. cvar_name
  663. for cvar_name, cvar in cls.computed_vars.items()
  664. if not cvar._cache
  665. )
  666. # Any substate containing a ComputedVar with cache=False always needs to be recomputed
  667. if cls._always_dirty_computed_vars:
  668. # Tell parent classes that this substate has always dirty computed vars
  669. state_name = cls.get_name()
  670. parent_state = cls.get_parent_state()
  671. while parent_state is not None:
  672. parent_state._always_dirty_substates.add(state_name)
  673. state_name, parent_state = (
  674. parent_state.get_name(),
  675. parent_state.get_parent_state(),
  676. )
  677. # Reset cached schema value
  678. cls._to_schema.cache_clear()
  679. @classmethod
  680. def _check_overridden_methods(cls):
  681. """Check for shadow methods and raise error if any.
  682. Raises:
  683. EventHandlerShadowsBuiltInStateMethod: When an event handler shadows an inbuilt state method.
  684. """
  685. overridden_methods = set()
  686. state_base_functions = cls._get_base_functions()
  687. for name, method in inspect.getmembers(cls, inspect.isfunction):
  688. # Check if the method is overridden and not a dunder method
  689. if (
  690. not name.startswith("__")
  691. and method.__name__ in state_base_functions
  692. and state_base_functions[method.__name__] != method
  693. ):
  694. overridden_methods.add(method.__name__)
  695. for method_name in overridden_methods:
  696. raise EventHandlerShadowsBuiltInStateMethod(
  697. f"The event handler name `{method_name}` shadows a builtin State method; use a different name instead"
  698. )
  699. @classmethod
  700. def _check_overridden_basevars(cls):
  701. """Check for shadow base vars and raise error if any.
  702. Raises:
  703. ComputedVarShadowsBaseVars: When a computed var shadows a base var.
  704. """
  705. for computed_var_ in cls._get_computed_vars():
  706. if computed_var_._js_expr in cls.__annotations__:
  707. raise ComputedVarShadowsBaseVars(
  708. f"The computed var name `{computed_var_._js_expr}` shadows a base var in {cls.__module__}.{cls.__name__}; use a different name instead"
  709. )
  710. @classmethod
  711. def _check_overridden_computed_vars(cls) -> None:
  712. """Check for shadow computed vars and raise error if any.
  713. Raises:
  714. ComputedVarShadowsStateVar: When a computed var shadows another.
  715. """
  716. for name, cv in cls.__dict__.items():
  717. if not is_computed_var(cv):
  718. continue
  719. name = cv._js_expr
  720. if name in cls.inherited_vars or name in cls.inherited_backend_vars:
  721. raise ComputedVarShadowsStateVar(
  722. f"The computed var name `{cv._js_expr}` shadows a var in {cls.__module__}.{cls.__name__}; use a different name instead"
  723. )
  724. @classmethod
  725. def get_skip_vars(cls) -> set[str]:
  726. """Get the vars to skip when serializing.
  727. Returns:
  728. The vars to skip when serializing.
  729. """
  730. return (
  731. set(cls.inherited_vars)
  732. | {
  733. "parent_state",
  734. "substates",
  735. "dirty_vars",
  736. "dirty_substates",
  737. "router_data",
  738. }
  739. | types.RESERVED_BACKEND_VAR_NAMES
  740. )
  741. @classmethod
  742. @functools.lru_cache()
  743. def get_parent_state(cls) -> Type[BaseState] | None:
  744. """Get the parent state.
  745. Raises:
  746. ValueError: If more than one parent state is found.
  747. Returns:
  748. The parent state.
  749. """
  750. parent_states = [
  751. base
  752. for base in cls.__bases__
  753. if issubclass(base, BaseState) and base is not BaseState and not base._mixin
  754. ]
  755. if len(parent_states) >= 2:
  756. raise ValueError(f"Only one parent state is allowed {parent_states}.")
  757. return parent_states[0] if len(parent_states) == 1 else None # type: ignore
  758. @classmethod
  759. def get_substates(cls) -> set[Type[BaseState]]:
  760. """Get the substates of the state.
  761. Returns:
  762. The substates of the state.
  763. """
  764. return cls.class_subclasses
  765. @classmethod
  766. @functools.lru_cache()
  767. def get_name(cls) -> str:
  768. """Get the name of the state.
  769. Returns:
  770. The name of the state.
  771. """
  772. module = cls.__module__.replace(".", "___")
  773. state_name = format.to_snake_case(f"{module}___{cls.__name__}")
  774. if environment.REFLEX_MINIFY_STATES.get():
  775. return get_minified_state_name(state_name)
  776. return state_name
  777. @classmethod
  778. @functools.lru_cache()
  779. def get_full_name(cls) -> str:
  780. """Get the full name of the state.
  781. Returns:
  782. The full name of the state.
  783. """
  784. name = cls.get_name()
  785. parent_state = cls.get_parent_state()
  786. if parent_state is not None:
  787. name = ".".join((parent_state.get_full_name(), name))
  788. return name
  789. @classmethod
  790. @functools.lru_cache()
  791. def get_class_substate(cls, path: Sequence[str] | str) -> Type[BaseState]:
  792. """Get the class substate.
  793. Args:
  794. path: The path to the substate.
  795. Returns:
  796. The class substate.
  797. Raises:
  798. ValueError: If the substate is not found.
  799. """
  800. if isinstance(path, str):
  801. path = tuple(path.split("."))
  802. if len(path) == 0:
  803. return cls
  804. if path[0] == cls.get_name():
  805. if len(path) == 1:
  806. return cls
  807. path = path[1:]
  808. for substate in cls.get_substates():
  809. if path[0] == substate.get_name():
  810. return substate.get_class_substate(path[1:])
  811. raise ValueError(f"Invalid path: {path}")
  812. @classmethod
  813. def get_class_var(cls, path: Sequence[str]) -> Any:
  814. """Get the class var.
  815. Args:
  816. path: The path to the var.
  817. Returns:
  818. The class var.
  819. Raises:
  820. ValueError: If the path is invalid.
  821. """
  822. path, name = path[:-1], path[-1]
  823. substate = cls.get_class_substate(tuple(path))
  824. if not hasattr(substate, name):
  825. raise ValueError(f"Invalid path: {path}")
  826. return getattr(substate, name)
  827. @classmethod
  828. def _init_var(cls, prop: Var):
  829. """Initialize a variable.
  830. Args:
  831. prop: The variable to initialize
  832. Raises:
  833. VarTypeError: if the variable has an incorrect type
  834. """
  835. from reflex.utils.exceptions import VarTypeError
  836. if not types.is_valid_var_type(prop._var_type):
  837. raise VarTypeError(
  838. "State vars must be primitive Python types, "
  839. "Plotly figures, Pandas dataframes, "
  840. "or subclasses of rx.Base. "
  841. f'Found var "{prop._js_expr}" with type {prop._var_type}.'
  842. )
  843. cls._set_var(prop)
  844. cls._create_setter(prop)
  845. cls._set_default_value(prop)
  846. @classmethod
  847. def add_var(cls, name: str, type_: Any, default_value: Any = None):
  848. """Add dynamically a variable to the State.
  849. The variable added this way can be used in the same way as a variable
  850. defined statically in the model.
  851. Args:
  852. name: The name of the variable
  853. type_: The type of the variable
  854. default_value: The default value of the variable
  855. Raises:
  856. NameError: if a variable of this name already exists
  857. """
  858. if name in cls.__fields__:
  859. raise NameError(
  860. f"The variable '{name}' already exist. Use a different name"
  861. )
  862. # create the variable based on name and type
  863. var = Var(
  864. _js_expr=format.format_state_name(cls.get_full_name()) + "." + name,
  865. _var_type=type_,
  866. _var_data=VarData.from_state(cls, name),
  867. ).guess_type()
  868. # add the pydantic field dynamically (must be done before _init_var)
  869. cls.add_field(var, default_value)
  870. cls._init_var(var)
  871. # update the internal dicts so the new variable is correctly handled
  872. cls.base_vars.update({name: var})
  873. cls.vars.update({name: var})
  874. # let substates know about the new variable
  875. for substate_class in cls.class_subclasses:
  876. substate_class.vars.setdefault(name, var)
  877. # Reinitialize dependency tracking dicts.
  878. cls._init_var_dependency_dicts()
  879. @classmethod
  880. def _set_var(cls, prop: Var):
  881. """Set the var as a class member.
  882. Args:
  883. prop: The var instance to set.
  884. """
  885. setattr(cls, prop._var_field_name, prop)
  886. @classmethod
  887. def _create_event_handler(cls, fn):
  888. """Create an event handler for the given function.
  889. Args:
  890. fn: The function to create an event handler for.
  891. Returns:
  892. The event handler.
  893. """
  894. return EventHandler(fn=fn, state_full_name=cls.get_full_name())
  895. @classmethod
  896. def _create_setvar(cls):
  897. """Create the setvar method for the state."""
  898. cls.setvar = cls.event_handlers["setvar"] = EventHandlerSetVar(state_cls=cls)
  899. @classmethod
  900. def _create_setter(cls, prop: Var):
  901. """Create a setter for the var.
  902. Args:
  903. prop: The var to create a setter for.
  904. """
  905. setter_name = prop._get_setter_name(include_state=False)
  906. if setter_name not in cls.__dict__:
  907. event_handler = cls._create_event_handler(prop._get_setter())
  908. cls.event_handlers[setter_name] = event_handler
  909. setattr(cls, setter_name, event_handler)
  910. @classmethod
  911. def _set_default_value(cls, prop: Var):
  912. """Set the default value for the var.
  913. Args:
  914. prop: The var to set the default value for.
  915. """
  916. # Get the pydantic field for the var.
  917. field = cls.get_fields()[prop._var_field_name]
  918. if field.required:
  919. default_value = prop._get_default_value()
  920. if default_value is not None:
  921. field.required = False
  922. field.default = default_value
  923. if (
  924. not field.required
  925. and field.default is None
  926. and not types.is_optional(prop._var_type)
  927. ):
  928. # Ensure frontend uses null coalescing when accessing.
  929. object.__setattr__(prop, "_var_type", Optional[prop._var_type])
  930. @classmethod
  931. def _get_var_default(cls, name: str, annotation_value: Any) -> Any:
  932. """Get the default value of a (backend) var.
  933. Args:
  934. name: The name of the var.
  935. annotation_value: The annotation value of the var.
  936. Returns:
  937. The default value of the var or None.
  938. """
  939. try:
  940. return getattr(cls, name)
  941. except AttributeError:
  942. try:
  943. return Var("", _var_type=annotation_value)._get_default_value()
  944. except TypeError:
  945. pass
  946. return None
  947. @staticmethod
  948. def _get_base_functions() -> dict[str, FunctionType]:
  949. """Get all functions of the state class excluding dunder methods.
  950. Returns:
  951. The functions of rx.State class as a dict.
  952. """
  953. return {
  954. func[0]: func[1]
  955. for func in inspect.getmembers(BaseState, predicate=inspect.isfunction)
  956. if not func[0].startswith("__")
  957. }
  958. @classmethod
  959. def _update_substate_inherited_vars(cls, vars_to_add: dict[str, Var]):
  960. """Update the inherited vars of substates recursively when new vars are added.
  961. Also updates the var dependency tracking dicts after adding vars.
  962. Args:
  963. vars_to_add: names to Var instances to add to substates
  964. """
  965. for substate_class in cls.class_subclasses:
  966. for name, var in vars_to_add.items():
  967. if types.is_backend_base_variable(name, cls):
  968. substate_class.backend_vars.setdefault(name, var)
  969. substate_class.inherited_backend_vars.setdefault(name, var)
  970. else:
  971. substate_class.vars.setdefault(name, var)
  972. substate_class.inherited_vars.setdefault(name, var)
  973. substate_class._update_substate_inherited_vars(vars_to_add)
  974. # Reinitialize dependency tracking dicts.
  975. cls._init_var_dependency_dicts()
  976. @classmethod
  977. def setup_dynamic_args(cls, args: dict[str, str]):
  978. """Set up args for easy access in renderer.
  979. Args:
  980. args: a dict of args
  981. """
  982. if not args:
  983. return
  984. cls._check_overwritten_dynamic_args(list(args.keys()))
  985. def argsingle_factory(param):
  986. def inner_func(self) -> str:
  987. return self.router.page.params.get(param, "")
  988. return inner_func
  989. def arglist_factory(param):
  990. def inner_func(self) -> List[str]:
  991. return self.router.page.params.get(param, [])
  992. return inner_func
  993. dynamic_vars = {}
  994. for param, value in args.items():
  995. if value == constants.RouteArgType.SINGLE:
  996. func = argsingle_factory(param)
  997. elif value == constants.RouteArgType.LIST:
  998. func = arglist_factory(param)
  999. else:
  1000. continue
  1001. dynamic_vars[param] = DynamicRouteVar(
  1002. fget=func,
  1003. cache=True,
  1004. _js_expr=param,
  1005. _var_data=VarData.from_state(cls),
  1006. )
  1007. setattr(cls, param, dynamic_vars[param])
  1008. # Update tracking dicts.
  1009. cls.computed_vars.update(dynamic_vars)
  1010. cls.vars.update(dynamic_vars)
  1011. cls._update_substate_inherited_vars(dynamic_vars)
  1012. @classmethod
  1013. def _check_overwritten_dynamic_args(cls, args: list[str]):
  1014. """Check if dynamic args are shadowing existing vars. Recursively checks all child states.
  1015. Args:
  1016. args: a dict of args
  1017. Raises:
  1018. DynamicRouteArgShadowsStateVar: If a dynamic arg is shadowing an existing var.
  1019. """
  1020. for arg in args:
  1021. if (
  1022. arg in cls.computed_vars
  1023. and not isinstance(cls.computed_vars[arg], DynamicRouteVar)
  1024. ) or arg in cls.base_vars:
  1025. raise DynamicRouteArgShadowsStateVar(
  1026. f"Dynamic route arg '{arg}' is shadowing an existing var in {cls.__module__}.{cls.__name__}"
  1027. )
  1028. for substate in cls.get_substates():
  1029. substate._check_overwritten_dynamic_args(args)
  1030. def __getattribute__(self, name: str) -> Any:
  1031. """Get the state var.
  1032. If the var is inherited, get the var from the parent state.
  1033. Args:
  1034. name: The name of the var.
  1035. Returns:
  1036. The value of the var.
  1037. """
  1038. # If the state hasn't been initialized yet, return the default value.
  1039. if not super().__getattribute__("__dict__"):
  1040. return super().__getattribute__(name)
  1041. inherited_vars = {
  1042. **super().__getattribute__("inherited_vars"),
  1043. **super().__getattribute__("inherited_backend_vars"),
  1044. }
  1045. # For now, handle router_data updates as a special case.
  1046. if name in inherited_vars or name == constants.ROUTER_DATA:
  1047. parent_state = super().__getattribute__("parent_state")
  1048. if parent_state is not None:
  1049. return getattr(parent_state, name)
  1050. # Allow event handlers to be called on the instance directly.
  1051. event_handlers = super().__getattribute__("event_handlers")
  1052. if name in event_handlers:
  1053. handler = event_handlers[name]
  1054. if handler.is_background:
  1055. fn = _no_chain_background_task(type(self), name, handler.fn)
  1056. else:
  1057. fn = functools.partial(handler.fn, self)
  1058. fn.__module__ = handler.fn.__module__ # type: ignore
  1059. fn.__qualname__ = handler.fn.__qualname__ # type: ignore
  1060. return fn
  1061. backend_vars = super().__getattribute__("_backend_vars")
  1062. if name in backend_vars:
  1063. value = backend_vars[name]
  1064. else:
  1065. value = super().__getattribute__(name)
  1066. if isinstance(value, EventHandler):
  1067. # The event handler is inherited from a parent, so let the parent convert
  1068. # it to a callable function.
  1069. parent_state = super().__getattribute__("parent_state")
  1070. if parent_state is not None:
  1071. return getattr(parent_state, name)
  1072. if MutableProxy._is_mutable_type(value) and (
  1073. name in super().__getattribute__("base_vars") or name in backend_vars
  1074. ):
  1075. # track changes in mutable containers (list, dict, set, etc)
  1076. return MutableProxy(wrapped=value, state=self, field_name=name)
  1077. return value
  1078. def __setattr__(self, name: str, value: Any):
  1079. """Set the attribute.
  1080. If the attribute is inherited, set the attribute on the parent state.
  1081. Args:
  1082. name: The name of the attribute.
  1083. value: The value of the attribute.
  1084. Raises:
  1085. SetUndefinedStateVarError: If a value of a var is set without first defining it.
  1086. """
  1087. if isinstance(value, MutableProxy):
  1088. # unwrap proxy objects when assigning back to the state
  1089. value = value.__wrapped__
  1090. # Set the var on the parent state.
  1091. inherited_vars = {**self.inherited_vars, **self.inherited_backend_vars}
  1092. if name in inherited_vars:
  1093. setattr(self.parent_state, name, value)
  1094. return
  1095. if name in self.backend_vars:
  1096. self._backend_vars.__setitem__(name, value)
  1097. self.dirty_vars.add(name)
  1098. self._mark_dirty()
  1099. return
  1100. if (
  1101. name not in self.vars
  1102. and name not in self.get_skip_vars()
  1103. and not name.startswith("__")
  1104. and not name.startswith(
  1105. f"_{getattr(type(self), '__original_name__', type(self).__name__)}__"
  1106. )
  1107. ):
  1108. raise SetUndefinedStateVarError(
  1109. f"The state variable '{name}' has not been defined in '{type(self).__name__}'. "
  1110. f"All state variables must be declared before they can be set."
  1111. )
  1112. fields = self.get_fields()
  1113. if name in fields:
  1114. field = fields[name]
  1115. field_type = field.outer_type_
  1116. if field.allow_none:
  1117. field_type = Union[field_type, None]
  1118. if not _isinstance(value, field_type):
  1119. console.deprecate(
  1120. "mismatched-type-assignment",
  1121. f"Tried to assign value {value} of type {type(value)} to field {type(self).__name__}.{name} of type {field_type}."
  1122. " This might lead to unexpected behavior.",
  1123. "0.6.5",
  1124. "0.7.0",
  1125. )
  1126. # Set the attribute.
  1127. super().__setattr__(name, value)
  1128. # Add the var to the dirty list.
  1129. if name in self.vars or name in self._computed_var_dependencies:
  1130. self.dirty_vars.add(name)
  1131. self._mark_dirty()
  1132. # For now, handle router_data updates as a special case
  1133. if name == constants.ROUTER_DATA:
  1134. self.dirty_vars.add(name)
  1135. self._mark_dirty()
  1136. def reset(self):
  1137. """Reset all the base vars to their default values."""
  1138. # Reset the base vars.
  1139. fields = self.get_fields()
  1140. for prop_name in self.base_vars:
  1141. if prop_name == constants.ROUTER:
  1142. continue # never reset the router data
  1143. field = fields[prop_name]
  1144. if default_factory := field.default_factory:
  1145. default = default_factory()
  1146. else:
  1147. default = copy.deepcopy(field.default)
  1148. setattr(self, prop_name, default)
  1149. # Reset the backend vars.
  1150. for prop_name, value in self.backend_vars.items():
  1151. setattr(self, prop_name, copy.deepcopy(value))
  1152. # Recursively reset the substates.
  1153. for substate in self.substates.values():
  1154. substate.reset()
  1155. def _reset_client_storage(self):
  1156. """Reset client storage base vars to their default values."""
  1157. # Client-side storage is reset during hydrate so that clearing cookies
  1158. # on the browser also resets the values on the backend.
  1159. fields = self.get_fields()
  1160. for prop_name in self.base_vars:
  1161. field = fields[prop_name]
  1162. if isinstance(field.default, ClientStorageBase) or (
  1163. isinstance(field.type_, type)
  1164. and issubclass(field.type_, ClientStorageBase)
  1165. ):
  1166. setattr(self, prop_name, copy.deepcopy(field.default))
  1167. # Recursively reset the substate client storage.
  1168. for substate in self.substates.values():
  1169. substate._reset_client_storage()
  1170. def get_substate(self, path: Sequence[str]) -> BaseState:
  1171. """Get the substate.
  1172. Args:
  1173. path: The path to the substate.
  1174. Returns:
  1175. The substate.
  1176. Raises:
  1177. ValueError: If the substate is not found.
  1178. """
  1179. if len(path) == 0:
  1180. return self
  1181. if path[0] == self.get_name():
  1182. if len(path) == 1:
  1183. return self
  1184. path = path[1:]
  1185. if path[0] not in self.substates:
  1186. raise ValueError(f"Invalid path: {path}")
  1187. return self.substates[path[0]].get_substate(path[1:])
  1188. @classmethod
  1189. def _get_common_ancestor(cls, other: Type[BaseState]) -> str:
  1190. """Find the name of the nearest common ancestor shared by this and the other state.
  1191. Args:
  1192. other: The other state.
  1193. Returns:
  1194. Full name of the nearest common ancestor.
  1195. """
  1196. common_ancestor_parts = []
  1197. for part1, part2 in zip(
  1198. cls.get_full_name().split("."),
  1199. other.get_full_name().split("."),
  1200. ):
  1201. if part1 != part2:
  1202. break
  1203. common_ancestor_parts.append(part1)
  1204. return ".".join(common_ancestor_parts)
  1205. @classmethod
  1206. def _determine_missing_parent_states(
  1207. cls, target_state_cls: Type[BaseState]
  1208. ) -> tuple[str, list[str]]:
  1209. """Determine the missing parent states between the target_state_cls and common ancestor of this state.
  1210. Args:
  1211. target_state_cls: The class of the state to find missing parent states for.
  1212. Returns:
  1213. The name of the common ancestor and the list of missing parent states.
  1214. """
  1215. common_ancestor_name = cls._get_common_ancestor(target_state_cls)
  1216. common_ancestor_parts = common_ancestor_name.split(".")
  1217. target_state_parts = tuple(target_state_cls.get_full_name().split("."))
  1218. relative_target_state_parts = target_state_parts[len(common_ancestor_parts) :]
  1219. # Determine which parent states to fetch from the common ancestor down to the target_state_cls.
  1220. fetch_parent_states = [common_ancestor_name]
  1221. for relative_parent_state_name in relative_target_state_parts:
  1222. fetch_parent_states.append(
  1223. ".".join((fetch_parent_states[-1], relative_parent_state_name))
  1224. )
  1225. return common_ancestor_name, fetch_parent_states[1:-1]
  1226. def _get_parent_states(self) -> list[tuple[str, BaseState]]:
  1227. """Get all parent state instances up to the root of the state tree.
  1228. Returns:
  1229. A list of tuples containing the name and the instance of each parent state.
  1230. """
  1231. parent_states_with_name = []
  1232. parent_state = self
  1233. while parent_state.parent_state is not None:
  1234. parent_state = parent_state.parent_state
  1235. parent_states_with_name.append((parent_state.get_full_name(), parent_state))
  1236. return parent_states_with_name
  1237. def _get_root_state(self) -> BaseState:
  1238. """Get the root state of the state tree.
  1239. Returns:
  1240. The root state of the state tree.
  1241. """
  1242. parent_state = self
  1243. while parent_state.parent_state is not None:
  1244. parent_state = parent_state.parent_state
  1245. return parent_state
  1246. async def _populate_parent_states(self, target_state_cls: Type[BaseState]):
  1247. """Populate substates in the tree between the target_state_cls and common ancestor of this state.
  1248. Args:
  1249. target_state_cls: The class of the state to populate parent states for.
  1250. Returns:
  1251. The parent state instance of target_state_cls.
  1252. Raises:
  1253. RuntimeError: If redis is not used in this backend process.
  1254. """
  1255. state_manager = get_state_manager()
  1256. if not isinstance(state_manager, StateManagerRedis):
  1257. raise RuntimeError(
  1258. f"Cannot populate parent states of {target_state_cls.get_full_name()} without redis. "
  1259. "(All states should already be available -- this is likely a bug).",
  1260. )
  1261. # Find the missing parent states up to the common ancestor.
  1262. (
  1263. common_ancestor_name,
  1264. missing_parent_states,
  1265. ) = self._determine_missing_parent_states(target_state_cls)
  1266. # Fetch all missing parent states and link them up to the common ancestor.
  1267. parent_states_tuple = self._get_parent_states()
  1268. root_state = parent_states_tuple[-1][1]
  1269. parent_states_by_name = dict(parent_states_tuple)
  1270. parent_state = parent_states_by_name[common_ancestor_name]
  1271. for parent_state_name in missing_parent_states:
  1272. try:
  1273. parent_state = root_state.get_substate(parent_state_name.split("."))
  1274. # The requested state is already cached, do NOT fetch it again.
  1275. continue
  1276. except ValueError:
  1277. # The requested state is missing, fetch from redis.
  1278. pass
  1279. parent_state = await state_manager.get_state(
  1280. token=_substate_key(
  1281. self.router.session.client_token, parent_state_name
  1282. ),
  1283. top_level=False,
  1284. get_substates=False,
  1285. parent_state=parent_state,
  1286. )
  1287. # Return the direct parent of target_state_cls for subsequent linking.
  1288. return parent_state
  1289. def _get_state_from_cache(self, state_cls: Type[BaseState]) -> BaseState:
  1290. """Get a state instance from the cache.
  1291. Args:
  1292. state_cls: The class of the state.
  1293. Returns:
  1294. The instance of state_cls associated with this state's client_token.
  1295. """
  1296. root_state = self._get_root_state()
  1297. return root_state.get_substate(state_cls.get_full_name().split("."))
  1298. async def _get_state_from_redis(self, state_cls: Type[BaseState]) -> BaseState:
  1299. """Get a state instance from redis.
  1300. Args:
  1301. state_cls: The class of the state.
  1302. Returns:
  1303. The instance of state_cls associated with this state's client_token.
  1304. Raises:
  1305. RuntimeError: If redis is not used in this backend process.
  1306. """
  1307. # Fetch all missing parent states from redis.
  1308. parent_state_of_state_cls = await self._populate_parent_states(state_cls)
  1309. # Then get the target state and all its substates.
  1310. state_manager = get_state_manager()
  1311. if not isinstance(state_manager, StateManagerRedis):
  1312. raise RuntimeError(
  1313. f"Requested state {state_cls.get_full_name()} is not cached and cannot be accessed without redis. "
  1314. "(All states should already be available -- this is likely a bug).",
  1315. )
  1316. return await state_manager.get_state(
  1317. token=_substate_key(self.router.session.client_token, state_cls),
  1318. top_level=False,
  1319. get_substates=True,
  1320. parent_state=parent_state_of_state_cls,
  1321. )
  1322. async def get_state(self, state_cls: Type[BaseState]) -> BaseState:
  1323. """Get an instance of the state associated with this token.
  1324. Allows for arbitrary access to sibling states from within an event handler.
  1325. Args:
  1326. state_cls: The class of the state.
  1327. Returns:
  1328. The instance of state_cls associated with this state's client_token.
  1329. """
  1330. # Fast case - if this state instance is already cached, get_substate from root state.
  1331. try:
  1332. return self._get_state_from_cache(state_cls)
  1333. except ValueError:
  1334. pass
  1335. # Slow case - fetch missing parent states from redis.
  1336. return await self._get_state_from_redis(state_cls)
  1337. def _get_event_handler(
  1338. self, event: Event
  1339. ) -> tuple[BaseState | StateProxy, EventHandler]:
  1340. """Get the event handler for the given event.
  1341. Args:
  1342. event: The event to get the handler for.
  1343. Returns:
  1344. The event handler.
  1345. Raises:
  1346. ValueError: If the event handler or substate is not found.
  1347. """
  1348. # Get the event handler.
  1349. path = event.name.split(".")
  1350. path, name = path[:-1], path[-1]
  1351. substate = self.get_substate(path)
  1352. if not substate:
  1353. raise ValueError(
  1354. "The value of state cannot be None when processing an event."
  1355. )
  1356. handler = substate.event_handlers[name]
  1357. # For background tasks, proxy the state
  1358. if handler.is_background:
  1359. substate = StateProxy(substate)
  1360. return substate, handler
  1361. async def _process(self, event: Event) -> AsyncIterator[StateUpdate]:
  1362. """Obtain event info and process event.
  1363. Args:
  1364. event: The event to process.
  1365. Yields:
  1366. The state update after processing the event.
  1367. """
  1368. # Get the event handler.
  1369. substate, handler = self._get_event_handler(event)
  1370. # Run the event generator and yield state updates.
  1371. async for update in self._process_event(
  1372. handler=handler,
  1373. state=substate,
  1374. payload=event.payload,
  1375. ):
  1376. yield update
  1377. def _check_valid(self, handler: EventHandler, events: Any) -> Any:
  1378. """Check if the events yielded are valid. They must be EventHandlers or EventSpecs.
  1379. Args:
  1380. handler: EventHandler.
  1381. events: The events to be checked.
  1382. Raises:
  1383. TypeError: If any of the events are not valid.
  1384. Returns:
  1385. The events as they are if valid.
  1386. """
  1387. def _is_valid_type(events: Any) -> bool:
  1388. return isinstance(events, (Event, EventHandler, EventSpec))
  1389. if events is None or _is_valid_type(events):
  1390. return events
  1391. try:
  1392. if all(_is_valid_type(e) for e in events):
  1393. return events
  1394. except TypeError:
  1395. pass
  1396. raise TypeError(
  1397. f"Your handler {handler.fn.__qualname__} must only return/yield: None, Events or other EventHandlers referenced by their class (not using `self`)"
  1398. )
  1399. def _as_state_update(
  1400. self,
  1401. handler: EventHandler,
  1402. events: EventSpec | list[EventSpec] | None,
  1403. final: bool,
  1404. ) -> StateUpdate:
  1405. """Convert the events to a StateUpdate.
  1406. Fixes the events and checks for validity before converting.
  1407. Args:
  1408. handler: The handler where the events originated from.
  1409. events: The events to queue with the update.
  1410. final: Whether the handler is done processing.
  1411. Returns:
  1412. The valid StateUpdate containing the events and final flag.
  1413. """
  1414. # get the delta from the root of the state tree
  1415. state = self._get_root_state()
  1416. token = self.router.session.client_token
  1417. # Convert valid EventHandler and EventSpec into Event
  1418. fixed_events = fix_events(self._check_valid(handler, events), token)
  1419. try:
  1420. # Get the delta after processing the event.
  1421. delta = state.get_delta()
  1422. state._clean()
  1423. return StateUpdate(
  1424. delta=delta,
  1425. events=fixed_events,
  1426. final=final if not handler.is_background else True,
  1427. )
  1428. except Exception as ex:
  1429. state._clean()
  1430. app_instance = getattr(prerequisites.get_app(), constants.CompileVars.APP)
  1431. event_specs = app_instance.backend_exception_handler(ex)
  1432. if event_specs is None:
  1433. return StateUpdate()
  1434. event_specs_correct_type = cast(
  1435. Union[List[Union[EventSpec, EventHandler]], None],
  1436. [event_specs] if isinstance(event_specs, EventSpec) else event_specs,
  1437. )
  1438. fixed_events = fix_events(
  1439. event_specs_correct_type,
  1440. token,
  1441. router_data=state.router_data,
  1442. )
  1443. return StateUpdate(
  1444. events=fixed_events,
  1445. final=True,
  1446. )
  1447. async def _process_event(
  1448. self, handler: EventHandler, state: BaseState | StateProxy, payload: Dict
  1449. ) -> AsyncIterator[StateUpdate]:
  1450. """Process event.
  1451. Args:
  1452. handler: EventHandler to process.
  1453. state: State to process the handler.
  1454. payload: The event payload.
  1455. Yields:
  1456. StateUpdate object
  1457. """
  1458. from reflex.utils import telemetry
  1459. # Get the function to process the event.
  1460. fn = functools.partial(handler.fn, state)
  1461. try:
  1462. type_hints = typing.get_type_hints(handler.fn)
  1463. except Exception:
  1464. type_hints = {}
  1465. for arg, value in list(payload.items()):
  1466. hinted_args = type_hints.get(arg, Any)
  1467. if hinted_args is Any:
  1468. continue
  1469. if is_union(hinted_args):
  1470. if value is None:
  1471. continue
  1472. hinted_args = value_inside_optional(hinted_args)
  1473. if (
  1474. isinstance(value, dict)
  1475. and inspect.isclass(hinted_args)
  1476. and not types.is_generic_alias(hinted_args) # py3.9-py3.10
  1477. ):
  1478. if issubclass(hinted_args, Model):
  1479. # Remove non-fields from the payload
  1480. payload[arg] = hinted_args(
  1481. **{
  1482. key: value
  1483. for key, value in value.items()
  1484. if key in hinted_args.__fields__
  1485. }
  1486. )
  1487. elif dataclasses.is_dataclass(hinted_args) or issubclass(
  1488. hinted_args, (Base, BaseModelV1, BaseModelV2)
  1489. ):
  1490. payload[arg] = hinted_args(**value)
  1491. if isinstance(value, list) and (hinted_args is set or hinted_args is Set):
  1492. payload[arg] = set(value)
  1493. if isinstance(value, list) and (
  1494. hinted_args is tuple or hinted_args is Tuple
  1495. ):
  1496. payload[arg] = tuple(value)
  1497. # Wrap the function in a try/except block.
  1498. try:
  1499. # Handle async functions.
  1500. if asyncio.iscoroutinefunction(fn.func):
  1501. events = await fn(**payload)
  1502. # Handle regular functions.
  1503. else:
  1504. events = fn(**payload)
  1505. # Handle async generators.
  1506. if inspect.isasyncgen(events):
  1507. async for event in events:
  1508. yield state._as_state_update(handler, event, final=False)
  1509. yield state._as_state_update(handler, events=None, final=True)
  1510. # Handle regular generators.
  1511. elif inspect.isgenerator(events):
  1512. try:
  1513. while True:
  1514. yield state._as_state_update(handler, next(events), final=False)
  1515. except StopIteration as si:
  1516. # the "return" value of the generator is not available
  1517. # in the loop, we must catch StopIteration to access it
  1518. if si.value is not None:
  1519. yield state._as_state_update(handler, si.value, final=False)
  1520. yield state._as_state_update(handler, events=None, final=True)
  1521. # Handle regular event chains.
  1522. else:
  1523. yield state._as_state_update(handler, events, final=True)
  1524. # If an error occurs, throw a window alert.
  1525. except Exception as ex:
  1526. telemetry.send_error(ex, context="backend")
  1527. app_instance = getattr(prerequisites.get_app(), constants.CompileVars.APP)
  1528. event_specs = app_instance.backend_exception_handler(ex)
  1529. yield state._as_state_update(
  1530. handler,
  1531. event_specs,
  1532. final=True,
  1533. )
  1534. def _mark_dirty_computed_vars(self) -> None:
  1535. """Mark ComputedVars that need to be recalculated based on dirty_vars."""
  1536. dirty_vars = self.dirty_vars
  1537. while dirty_vars:
  1538. calc_vars, dirty_vars = dirty_vars, set()
  1539. for cvar in self._dirty_computed_vars(from_vars=calc_vars):
  1540. self.dirty_vars.add(cvar)
  1541. dirty_vars.add(cvar)
  1542. actual_var = self.computed_vars.get(cvar)
  1543. if actual_var is not None:
  1544. actual_var.mark_dirty(instance=self)
  1545. def _expired_computed_vars(self) -> set[str]:
  1546. """Determine ComputedVars that need to be recalculated based on the expiration time.
  1547. Returns:
  1548. Set of computed vars to include in the delta.
  1549. """
  1550. return set(
  1551. cvar
  1552. for cvar in self.computed_vars
  1553. if self.computed_vars[cvar].needs_update(instance=self)
  1554. )
  1555. def _dirty_computed_vars(
  1556. self, from_vars: set[str] | None = None, include_backend: bool = True
  1557. ) -> set[str]:
  1558. """Determine ComputedVars that need to be recalculated based on the given vars.
  1559. Args:
  1560. from_vars: find ComputedVar that depend on this set of vars. If unspecified, will use the dirty_vars.
  1561. include_backend: whether to include backend vars in the calculation.
  1562. Returns:
  1563. Set of computed vars to include in the delta.
  1564. """
  1565. return set(
  1566. cvar
  1567. for dirty_var in from_vars or self.dirty_vars
  1568. for cvar in self._computed_var_dependencies[dirty_var]
  1569. if include_backend or not self.computed_vars[cvar]._backend
  1570. )
  1571. @classmethod
  1572. def _potentially_dirty_substates(cls) -> set[Type[BaseState]]:
  1573. """Determine substates which could be affected by dirty vars in this state.
  1574. Returns:
  1575. Set of State classes that may need to be fetched to recalc computed vars.
  1576. """
  1577. # _always_dirty_substates need to be fetched to recalc computed vars.
  1578. fetch_substates = set(
  1579. cls.get_class_substate((cls.get_name(), *substate_name.split(".")))
  1580. for substate_name in cls._always_dirty_substates
  1581. )
  1582. for dependent_substates in cls._substate_var_dependencies.values():
  1583. fetch_substates.update(
  1584. set(
  1585. cls.get_class_substate((cls.get_name(), *substate_name.split(".")))
  1586. for substate_name in dependent_substates
  1587. )
  1588. )
  1589. return fetch_substates
  1590. def get_delta(self) -> Delta:
  1591. """Get the delta for the state.
  1592. Returns:
  1593. The delta for the state.
  1594. """
  1595. delta = {}
  1596. # Apply dirty variables down into substates
  1597. self.dirty_vars.update(self._always_dirty_computed_vars)
  1598. self._mark_dirty()
  1599. frontend_computed_vars: set[str] = {
  1600. name for name, cv in self.computed_vars.items() if not cv._backend
  1601. }
  1602. # Return the dirty vars for this instance, any cached/dependent computed vars,
  1603. # and always dirty computed vars (cache=False)
  1604. delta_vars = (
  1605. self.dirty_vars.intersection(self.base_vars)
  1606. .union(self.dirty_vars.intersection(frontend_computed_vars))
  1607. .union(self._dirty_computed_vars(include_backend=False))
  1608. .union(self._always_dirty_computed_vars)
  1609. )
  1610. subdelta: Dict[str, Any] = {
  1611. prop: self.get_value(prop)
  1612. for prop in delta_vars
  1613. if not types.is_backend_base_variable(prop, type(self))
  1614. }
  1615. if len(subdelta) > 0:
  1616. delta[self.get_full_name()] = subdelta
  1617. # Recursively find the substate deltas.
  1618. substates = self.substates
  1619. for substate in self.dirty_substates.union(self._always_dirty_substates):
  1620. delta.update(substates[substate].get_delta())
  1621. # Return the delta.
  1622. return delta
  1623. def _mark_dirty(self):
  1624. """Mark the substate and all parent states as dirty."""
  1625. state_name = self.get_name()
  1626. if (
  1627. self.parent_state is not None
  1628. and state_name not in self.parent_state.dirty_substates
  1629. ):
  1630. self.parent_state.dirty_substates.add(self.get_name())
  1631. self.parent_state._mark_dirty()
  1632. # Append expired computed vars to dirty_vars to trigger recalculation
  1633. self.dirty_vars.update(self._expired_computed_vars())
  1634. # have to mark computed vars dirty to allow access to newly computed
  1635. # values within the same ComputedVar function
  1636. self._mark_dirty_computed_vars()
  1637. self._mark_dirty_substates()
  1638. def _mark_dirty_substates(self):
  1639. """Propagate dirty var / computed var status into substates."""
  1640. substates = self.substates
  1641. for var in self.dirty_vars:
  1642. for substate_name in self._substate_var_dependencies[var]:
  1643. self.dirty_substates.add(substate_name)
  1644. substate = substates[substate_name]
  1645. substate.dirty_vars.add(var)
  1646. substate._mark_dirty()
  1647. def _update_was_touched(self):
  1648. """Update the _was_touched flag based on dirty_vars."""
  1649. if self.dirty_vars and not self._was_touched:
  1650. for var in self.dirty_vars:
  1651. if var in self.base_vars or var in self._backend_vars:
  1652. self._was_touched = True
  1653. break
  1654. def _get_was_touched(self) -> bool:
  1655. """Check current dirty_vars and flag to determine if state instance was modified.
  1656. If any dirty vars belong to this state, mark _was_touched.
  1657. This flag determines whether this state instance should be persisted to redis.
  1658. Returns:
  1659. Whether this state instance was ever modified.
  1660. """
  1661. # Ensure the flag is up to date based on the current dirty_vars
  1662. self._update_was_touched()
  1663. return self._was_touched
  1664. def _clean(self):
  1665. """Reset the dirty vars."""
  1666. # Update touched status before cleaning dirty_vars.
  1667. self._update_was_touched()
  1668. # Recursively clean the substates.
  1669. for substate in self.dirty_substates:
  1670. if substate not in self.substates:
  1671. continue
  1672. self.substates[substate]._clean()
  1673. # Clean this state.
  1674. self.dirty_vars = set()
  1675. self.dirty_substates = set()
  1676. def get_value(self, key: str) -> Any:
  1677. """Get the value of a field (without proxying).
  1678. The returned value will NOT track dirty state updates.
  1679. Args:
  1680. key: The key of the field.
  1681. Returns:
  1682. The value of the field.
  1683. """
  1684. value = super().get_value(key)
  1685. if isinstance(value, MutableProxy):
  1686. return value.__wrapped__
  1687. return value
  1688. def dict(
  1689. self, include_computed: bool = True, initial: bool = False, **kwargs
  1690. ) -> dict[str, Any]:
  1691. """Convert the object to a dictionary.
  1692. Args:
  1693. include_computed: Whether to include computed vars.
  1694. initial: Whether to get the initial value of computed vars.
  1695. **kwargs: Kwargs to pass to the pydantic dict method.
  1696. Returns:
  1697. The object as a dictionary.
  1698. """
  1699. if include_computed:
  1700. # Apply dirty variables down into substates to allow never-cached ComputedVar to
  1701. # trigger recalculation of dependent vars
  1702. self.dirty_vars.update(self._always_dirty_computed_vars)
  1703. self._mark_dirty()
  1704. base_vars = {
  1705. prop_name: self.get_value(prop_name) for prop_name in self.base_vars
  1706. }
  1707. if initial and include_computed:
  1708. computed_vars = {
  1709. # Include initial computed vars.
  1710. prop_name: (
  1711. cv._initial_value
  1712. if is_computed_var(cv)
  1713. and not isinstance(cv._initial_value, types.Unset)
  1714. else self.get_value(prop_name)
  1715. )
  1716. for prop_name, cv in self.computed_vars.items()
  1717. if not cv._backend
  1718. }
  1719. elif include_computed:
  1720. computed_vars = {
  1721. # Include the computed vars.
  1722. prop_name: self.get_value(prop_name)
  1723. for prop_name, cv in self.computed_vars.items()
  1724. if not cv._backend
  1725. }
  1726. else:
  1727. computed_vars = {}
  1728. variables = {**base_vars, **computed_vars}
  1729. d = {
  1730. self.get_full_name(): {k: variables[k] for k in sorted(variables)},
  1731. }
  1732. for substate_d in [
  1733. v.dict(include_computed=include_computed, initial=initial, **kwargs)
  1734. for v in self.substates.values()
  1735. ]:
  1736. d.update(substate_d)
  1737. return d
  1738. async def __aenter__(self) -> BaseState:
  1739. """Enter the async context manager protocol.
  1740. This should not be used for the State class, but exists for
  1741. type-compatibility with StateProxy.
  1742. Raises:
  1743. TypeError: always, because async contextmanager protocol is only supported for background task.
  1744. """
  1745. raise TypeError(
  1746. "Only background task should use `async with self` to modify state."
  1747. )
  1748. async def __aexit__(self, *exc_info: Any) -> None:
  1749. """Exit the async context manager protocol.
  1750. This should not be used for the State class, but exists for
  1751. type-compatibility with StateProxy.
  1752. Args:
  1753. exc_info: The exception info tuple.
  1754. """
  1755. pass
  1756. def __getstate__(self):
  1757. """Get the state for redis serialization.
  1758. This method is called by pickle to serialize the object.
  1759. It explicitly removes parent_state and substates because those are serialized separately
  1760. by the StateManagerRedis to allow for better horizontal scaling as state size increases.
  1761. Returns:
  1762. The state dict for serialization.
  1763. """
  1764. state = super().__getstate__()
  1765. state["__dict__"] = state["__dict__"].copy()
  1766. if state["__dict__"].get("parent_state") is not None:
  1767. # Do not serialize router data in substates (only the root state).
  1768. state["__dict__"].pop("router", None)
  1769. state["__dict__"].pop("router_data", None)
  1770. # Never serialize parent_state or substates.
  1771. state["__dict__"]["parent_state"] = None
  1772. state["__dict__"]["substates"] = {}
  1773. state["__dict__"].pop("_was_touched", None)
  1774. # Remove all inherited vars.
  1775. for inherited_var_name in self.inherited_vars:
  1776. state["__dict__"].pop(inherited_var_name, None)
  1777. return state
  1778. def _check_state_size(
  1779. self,
  1780. pickle_state_size: int,
  1781. ):
  1782. """Print a warning when the state is too large.
  1783. Args:
  1784. pickle_state_size: The size of the pickled state.
  1785. Raises:
  1786. StateTooLargeError: If the state is too large.
  1787. """
  1788. state_full_name = self.get_full_name()
  1789. if (
  1790. state_full_name not in _WARNED_ABOUT_STATE_SIZE
  1791. and pickle_state_size > TOO_LARGE_SERIALIZED_STATE
  1792. and self.substates
  1793. ):
  1794. msg = (
  1795. f"State {state_full_name} serializes to {pickle_state_size} bytes "
  1796. + "which may present performance issues. Consider reducing the size of this state."
  1797. )
  1798. if environment.REFLEX_PERF_MODE.get() == PerformanceMode.WARN:
  1799. console.warn(msg)
  1800. elif environment.REFLEX_PERF_MODE.get() == PerformanceMode.RAISE:
  1801. raise StateTooLargeError(msg)
  1802. _WARNED_ABOUT_STATE_SIZE.add(state_full_name)
  1803. @classmethod
  1804. @functools.lru_cache()
  1805. def _to_schema(cls) -> str:
  1806. """Convert a state to a schema.
  1807. Returns:
  1808. The hash of the schema.
  1809. """
  1810. def _field_tuple(
  1811. field_name: str,
  1812. ) -> Tuple[str, str, Any, Union[bool, None], Any]:
  1813. model_field = cls.__fields__[field_name]
  1814. return (
  1815. field_name,
  1816. model_field.name,
  1817. _serialize_type(model_field.type_),
  1818. (
  1819. model_field.required
  1820. if isinstance(model_field.required, bool)
  1821. else None
  1822. ),
  1823. (model_field.default if is_serializable(model_field.default) else None),
  1824. )
  1825. return md5(
  1826. pickle.dumps(
  1827. list(sorted(_field_tuple(field_name) for field_name in cls.base_vars))
  1828. )
  1829. ).hexdigest()
  1830. def _serialize(self) -> bytes:
  1831. """Serialize the state for redis.
  1832. Returns:
  1833. The serialized state.
  1834. """
  1835. try:
  1836. pickle_state = pickle.dumps((self._to_schema(), self))
  1837. if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
  1838. self._check_state_size(len(pickle_state))
  1839. return pickle_state
  1840. except HANDLED_PICKLE_ERRORS as og_pickle_error:
  1841. error = (
  1842. f"Failed to serialize state {self.get_full_name()} due to unpicklable object. "
  1843. "This state will not be persisted. "
  1844. )
  1845. try:
  1846. import dill
  1847. return dill.dumps((self._to_schema(), self))
  1848. except ImportError:
  1849. error += (
  1850. f"Pickle error: {og_pickle_error}. "
  1851. "Consider `pip install 'dill>=0.3.8'` for more exotic serialization support."
  1852. )
  1853. except HANDLED_PICKLE_ERRORS as ex:
  1854. error += f"Dill was also unable to pickle the state: {ex}"
  1855. console.warn(error)
  1856. return b""
  1857. @classmethod
  1858. def _deserialize(
  1859. cls, data: bytes | None = None, fp: BinaryIO | None = None
  1860. ) -> BaseState:
  1861. """Deserialize the state from redis/disk.
  1862. data and fp are mutually exclusive, but one must be provided.
  1863. Args:
  1864. data: The serialized state data.
  1865. fp: The file pointer to the serialized state data.
  1866. Returns:
  1867. The deserialized state.
  1868. Raises:
  1869. ValueError: If both data and fp are provided, or neither are provided.
  1870. StateSchemaMismatchError: If the state schema does not match the expected schema.
  1871. """
  1872. if data is not None and fp is None:
  1873. (substate_schema, state) = pickle.loads(data)
  1874. elif fp is not None and data is None:
  1875. (substate_schema, state) = pickle.load(fp)
  1876. else:
  1877. raise ValueError("Only one of `data` or `fp` must be provided")
  1878. if substate_schema != state._to_schema():
  1879. raise StateSchemaMismatchError()
  1880. return state
  1881. class State(BaseState):
  1882. """The app Base State."""
  1883. # The hydrated bool.
  1884. is_hydrated: bool = False
  1885. T = TypeVar("T", bound=BaseState)
  1886. def dynamic(func: Callable[[T], Component]):
  1887. """Create a dynamically generated components from a state class.
  1888. Args:
  1889. func: The function to generate the component.
  1890. Returns:
  1891. The dynamically generated component.
  1892. Raises:
  1893. DynamicComponentInvalidSignature: If the function does not have exactly one parameter.
  1894. DynamicComponentInvalidSignature: If the function does not have a type hint for the state class.
  1895. """
  1896. number_of_parameters = len(inspect.signature(func).parameters)
  1897. func_signature = get_type_hints(func)
  1898. if "return" in func_signature:
  1899. func_signature.pop("return")
  1900. values = list(func_signature.values())
  1901. if number_of_parameters != 1:
  1902. raise DynamicComponentInvalidSignature(
  1903. "The function must have exactly one parameter, which is the state class."
  1904. )
  1905. if len(values) != 1:
  1906. raise DynamicComponentInvalidSignature(
  1907. "You must provide a type hint for the state class in the function."
  1908. )
  1909. state_class: Type[T] = values[0]
  1910. def wrapper() -> Component:
  1911. from reflex.components.base.fragment import fragment
  1912. return fragment(state_class._evaluate(lambda state: func(state)))
  1913. return wrapper
  1914. class FrontendEventExceptionState(State):
  1915. """Substate for handling frontend exceptions."""
  1916. @event
  1917. def handle_frontend_exception(self, stack: str, component_stack: str) -> None:
  1918. """Handle frontend exceptions.
  1919. If a frontend exception handler is provided, it will be called.
  1920. Otherwise, the default frontend exception handler will be called.
  1921. Args:
  1922. stack: The stack trace of the exception.
  1923. component_stack: The stack trace of the component where the exception occurred.
  1924. """
  1925. app_instance = getattr(prerequisites.get_app(), constants.CompileVars.APP)
  1926. app_instance.frontend_exception_handler(Exception(stack))
  1927. class UpdateVarsInternalState(State):
  1928. """Substate for handling internal state var updates."""
  1929. async def update_vars_internal(self, vars: dict[str, Any]) -> None:
  1930. """Apply updates to fully qualified state vars.
  1931. The keys in `vars` should be in the form of `{state.get_full_name()}.{var_name}`,
  1932. and each value will be set on the appropriate substate instance.
  1933. This function is primarily used to apply cookie and local storage
  1934. updates from the frontend to the appropriate substate.
  1935. Args:
  1936. vars: The fully qualified vars and values to update.
  1937. """
  1938. for var, value in vars.items():
  1939. state_name, _, var_name = var.rpartition(".")
  1940. var_state_cls = State.get_class_substate(state_name)
  1941. var_state = await self.get_state(var_state_cls)
  1942. setattr(var_state, var_name, value)
  1943. class OnLoadInternalState(State):
  1944. """Substate for handling on_load event enumeration.
  1945. This is a separate substate to avoid deserializing the entire state tree for every page navigation.
  1946. """
  1947. def on_load_internal(self) -> list[Event | EventSpec] | None:
  1948. """Queue on_load handlers for the current page.
  1949. Returns:
  1950. The list of events to queue for on load handling.
  1951. """
  1952. # Do not app._compile()! It should be already compiled by now.
  1953. app = getattr(prerequisites.get_app(), constants.CompileVars.APP)
  1954. load_events = app.get_load_events(self.router.page.path)
  1955. if not load_events:
  1956. self.is_hydrated = True
  1957. return # Fast path for navigation with no on_load events defined.
  1958. self.is_hydrated = False
  1959. return [
  1960. *fix_events(
  1961. load_events,
  1962. self.router.session.client_token,
  1963. router_data=self.router_data,
  1964. ),
  1965. State.set_is_hydrated(True), # type: ignore
  1966. ]
  1967. class ComponentState(State, mixin=True):
  1968. """Base class to allow for the creation of a state instance per component.
  1969. This allows for the bundling of UI and state logic into a single class,
  1970. where each instance has a separate instance of the state.
  1971. Subclass this class and define vars and event handlers in the traditional way.
  1972. Then define a `get_component` method that returns the UI for the component instance.
  1973. See the full [docs](https://reflex.dev/docs/substates/component-state/) for more.
  1974. Basic example:
  1975. ```python
  1976. # Subclass ComponentState and define vars and event handlers.
  1977. class Counter(rx.ComponentState):
  1978. # Define vars that change.
  1979. count: int = 0
  1980. # Define event handlers.
  1981. def increment(self):
  1982. self.count += 1
  1983. def decrement(self):
  1984. self.count -= 1
  1985. @classmethod
  1986. def get_component(cls, **props):
  1987. # Access the state vars and event handlers using `cls`.
  1988. return rx.hstack(
  1989. rx.button("Decrement", on_click=cls.decrement),
  1990. rx.text(cls.count),
  1991. rx.button("Increment", on_click=cls.increment),
  1992. **props,
  1993. )
  1994. counter = Counter.create()
  1995. ```
  1996. """
  1997. # The number of components created from this class.
  1998. _per_component_state_instance_count: ClassVar[int] = 0
  1999. def __init__(self, *args, **kwargs):
  2000. """Do not allow direct initialization of the ComponentState.
  2001. Args:
  2002. *args: The args to pass to the State init method.
  2003. **kwargs: The kwargs to pass to the State init method.
  2004. Raises:
  2005. ReflexRuntimeError: If the ComponentState is initialized directly.
  2006. """
  2007. if type(self)._mixin:
  2008. raise ReflexRuntimeError(
  2009. f"{ComponentState.__name__} {type(self).__name__} is not meant to be initialized directly. "
  2010. + "Use the `create` method to create a new instance and access the state via the `State` attribute."
  2011. )
  2012. super().__init__(*args, **kwargs)
  2013. @classmethod
  2014. def __init_subclass__(cls, mixin: bool = True, **kwargs):
  2015. """Overwrite mixin default to True.
  2016. Args:
  2017. mixin: Whether the subclass is a mixin and should not be initialized.
  2018. **kwargs: The kwargs to pass to the pydantic init_subclass method.
  2019. """
  2020. super().__init_subclass__(mixin=mixin, **kwargs)
  2021. @classmethod
  2022. def get_component(cls, *children, **props) -> "Component":
  2023. """Get the component instance.
  2024. Args:
  2025. children: The children of the component.
  2026. props: The props of the component.
  2027. Raises:
  2028. NotImplementedError: if the subclass does not override this method.
  2029. """
  2030. raise NotImplementedError(
  2031. f"{cls.__name__} must implement get_component to return the component instance."
  2032. )
  2033. @classmethod
  2034. def create(cls, *children, **props) -> "Component":
  2035. """Create a new instance of the Component.
  2036. Args:
  2037. children: The children of the component.
  2038. props: The props of the component.
  2039. Returns:
  2040. A new instance of the Component with an independent copy of the State.
  2041. """
  2042. cls._per_component_state_instance_count += 1
  2043. state_cls_name = f"{cls.__name__}_n{cls._per_component_state_instance_count}"
  2044. component_state = type(
  2045. state_cls_name,
  2046. (cls, State),
  2047. {"__module__": reflex.istate.dynamic.__name__},
  2048. mixin=False,
  2049. )
  2050. # Save a reference to the dynamic state for pickle/unpickle.
  2051. setattr(reflex.istate.dynamic, state_cls_name, component_state)
  2052. component = component_state.get_component(*children, **props)
  2053. component.State = component_state
  2054. return component
  2055. class StateProxy(wrapt.ObjectProxy):
  2056. """Proxy of a state instance to control mutability of vars for a background task.
  2057. Since a background task runs against a state instance without holding the
  2058. state_manager lock for the token, the reference may become stale if the same
  2059. state is modified by another event handler.
  2060. The proxy object ensures that writes to the state are blocked unless
  2061. explicitly entering a context which refreshes the state from state_manager
  2062. and holds the lock for the token until exiting the context. After exiting
  2063. the context, a StateUpdate may be emitted to the frontend to notify the
  2064. client of the state change.
  2065. A background task will be passed the `StateProxy` as `self`, so mutability
  2066. can be safely performed inside an `async with self` block.
  2067. class State(rx.State):
  2068. counter: int = 0
  2069. @rx.event(background=True)
  2070. async def bg_increment(self):
  2071. await asyncio.sleep(1)
  2072. async with self:
  2073. self.counter += 1
  2074. """
  2075. def __init__(
  2076. self, state_instance, parent_state_proxy: Optional["StateProxy"] = None
  2077. ):
  2078. """Create a proxy for a state instance.
  2079. If `get_state` is used on a StateProxy, the resulting state will be
  2080. linked to the given state via parent_state_proxy. The first state in the
  2081. chain is the state that initiated the background task.
  2082. Args:
  2083. state_instance: The state instance to proxy.
  2084. parent_state_proxy: The parent state proxy, for linked mutability and context tracking.
  2085. """
  2086. super().__init__(state_instance)
  2087. # compile is not relevant to backend logic
  2088. self._self_app = getattr(prerequisites.get_app(), constants.CompileVars.APP)
  2089. self._self_substate_path = tuple(state_instance.get_full_name().split("."))
  2090. self._self_actx = None
  2091. self._self_mutable = False
  2092. self._self_actx_lock = asyncio.Lock()
  2093. self._self_actx_lock_holder = None
  2094. self._self_parent_state_proxy = parent_state_proxy
  2095. def _is_mutable(self) -> bool:
  2096. """Check if the state is mutable.
  2097. Returns:
  2098. Whether the state is mutable.
  2099. """
  2100. if self._self_parent_state_proxy is not None:
  2101. return self._self_parent_state_proxy._is_mutable() or self._self_mutable
  2102. return self._self_mutable
  2103. async def __aenter__(self) -> StateProxy:
  2104. """Enter the async context manager protocol.
  2105. Sets mutability to True and enters the `App.modify_state` async context,
  2106. which refreshes the state from state_manager and holds the lock for the
  2107. given state token until exiting the context.
  2108. Background tasks should avoid blocking calls while inside the context.
  2109. Returns:
  2110. This StateProxy instance in mutable mode.
  2111. Raises:
  2112. ImmutableStateError: If the state is already mutable.
  2113. """
  2114. if self._self_parent_state_proxy is not None:
  2115. parent_state = (
  2116. await self._self_parent_state_proxy.__aenter__()
  2117. ).__wrapped__
  2118. super().__setattr__(
  2119. "__wrapped__",
  2120. await parent_state.get_state(
  2121. State.get_class_substate(self._self_substate_path)
  2122. ),
  2123. )
  2124. return self
  2125. current_task = asyncio.current_task()
  2126. if (
  2127. self._self_actx_lock.locked()
  2128. and current_task == self._self_actx_lock_holder
  2129. ):
  2130. raise ImmutableStateError(
  2131. "The state is already mutable. Do not nest `async with self` blocks."
  2132. )
  2133. await self._self_actx_lock.acquire()
  2134. self._self_actx_lock_holder = current_task
  2135. self._self_actx = self._self_app.modify_state(
  2136. token=_substate_key(
  2137. self.__wrapped__.router.session.client_token,
  2138. self._self_substate_path,
  2139. )
  2140. )
  2141. mutable_state = await self._self_actx.__aenter__()
  2142. super().__setattr__(
  2143. "__wrapped__", mutable_state.get_substate(self._self_substate_path)
  2144. )
  2145. self._self_mutable = True
  2146. return self
  2147. async def __aexit__(self, *exc_info: Any) -> None:
  2148. """Exit the async context manager protocol.
  2149. Sets proxy mutability to False and persists any state changes.
  2150. Args:
  2151. exc_info: The exception info tuple.
  2152. """
  2153. if self._self_parent_state_proxy is not None:
  2154. await self._self_parent_state_proxy.__aexit__(*exc_info)
  2155. return
  2156. if self._self_actx is None:
  2157. return
  2158. self._self_mutable = False
  2159. try:
  2160. await self._self_actx.__aexit__(*exc_info)
  2161. finally:
  2162. self._self_actx_lock_holder = None
  2163. self._self_actx_lock.release()
  2164. self._self_actx = None
  2165. def __enter__(self):
  2166. """Enter the regular context manager protocol.
  2167. This is not supported for background tasks, and exists only to raise a more useful exception
  2168. when the StateProxy is used incorrectly.
  2169. Raises:
  2170. TypeError: always, because only async contextmanager protocol is supported.
  2171. """
  2172. raise TypeError("Background task must use `async with self` to modify state.")
  2173. def __exit__(self, *exc_info: Any) -> None:
  2174. """Exit the regular context manager protocol.
  2175. Args:
  2176. exc_info: The exception info tuple.
  2177. """
  2178. pass
  2179. def __getattr__(self, name: str) -> Any:
  2180. """Get the attribute from the underlying state instance.
  2181. Args:
  2182. name: The name of the attribute.
  2183. Returns:
  2184. The value of the attribute.
  2185. Raises:
  2186. ImmutableStateError: If the state is not in mutable mode.
  2187. """
  2188. if name in ["substates", "parent_state"] and not self._is_mutable():
  2189. raise ImmutableStateError(
  2190. "Background task StateProxy is immutable outside of a context "
  2191. "manager. Use `async with self` to modify state."
  2192. )
  2193. value = super().__getattr__(name)
  2194. if not name.startswith("_self_") and isinstance(value, MutableProxy):
  2195. # ensure mutations to these containers are blocked unless proxy is _mutable
  2196. return ImmutableMutableProxy(
  2197. wrapped=value.__wrapped__,
  2198. state=self, # type: ignore
  2199. field_name=value._self_field_name,
  2200. )
  2201. if isinstance(value, functools.partial) and value.args[0] is self.__wrapped__:
  2202. # Rebind event handler to the proxy instance
  2203. value = functools.partial(
  2204. value.func,
  2205. self,
  2206. *value.args[1:],
  2207. **value.keywords,
  2208. )
  2209. if isinstance(value, MethodType) and value.__self__ is self.__wrapped__:
  2210. # Rebind methods to the proxy instance
  2211. value = type(value)(value.__func__, self) # type: ignore
  2212. return value
  2213. def __setattr__(self, name: str, value: Any) -> None:
  2214. """Set the attribute on the underlying state instance.
  2215. If the attribute is internal, set it on the proxy instance instead.
  2216. Args:
  2217. name: The name of the attribute.
  2218. value: The value of the attribute.
  2219. Raises:
  2220. ImmutableStateError: If the state is not in mutable mode.
  2221. """
  2222. if (
  2223. name.startswith("_self_") # wrapper attribute
  2224. or self._is_mutable() # lock held
  2225. # non-persisted state attribute
  2226. or name in self.__wrapped__.get_skip_vars()
  2227. ):
  2228. super().__setattr__(name, value)
  2229. return
  2230. raise ImmutableStateError(
  2231. "Background task StateProxy is immutable outside of a context "
  2232. "manager. Use `async with self` to modify state."
  2233. )
  2234. def get_substate(self, path: Sequence[str]) -> BaseState:
  2235. """Only allow substate access with lock held.
  2236. Args:
  2237. path: The path to the substate.
  2238. Returns:
  2239. The substate.
  2240. Raises:
  2241. ImmutableStateError: If the state is not in mutable mode.
  2242. """
  2243. if not self._is_mutable():
  2244. raise ImmutableStateError(
  2245. "Background task StateProxy is immutable outside of a context "
  2246. "manager. Use `async with self` to modify state."
  2247. )
  2248. return self.__wrapped__.get_substate(path)
  2249. async def get_state(self, state_cls: Type[BaseState]) -> BaseState:
  2250. """Get an instance of the state associated with this token.
  2251. Args:
  2252. state_cls: The class of the state.
  2253. Returns:
  2254. The state.
  2255. Raises:
  2256. ImmutableStateError: If the state is not in mutable mode.
  2257. """
  2258. if not self._is_mutable():
  2259. raise ImmutableStateError(
  2260. "Background task StateProxy is immutable outside of a context "
  2261. "manager. Use `async with self` to modify state."
  2262. )
  2263. return type(self)(
  2264. await self.__wrapped__.get_state(state_cls), parent_state_proxy=self
  2265. )
  2266. def _as_state_update(self, *args, **kwargs) -> StateUpdate:
  2267. """Temporarily allow mutability to access parent_state.
  2268. Args:
  2269. *args: The args to pass to the underlying state instance.
  2270. **kwargs: The kwargs to pass to the underlying state instance.
  2271. Returns:
  2272. The state update.
  2273. """
  2274. original_mutable = self._self_mutable
  2275. self._self_mutable = True
  2276. try:
  2277. return self.__wrapped__._as_state_update(*args, **kwargs)
  2278. finally:
  2279. self._self_mutable = original_mutable
  2280. @dataclasses.dataclass(
  2281. frozen=True,
  2282. )
  2283. class StateUpdate:
  2284. """A state update sent to the frontend."""
  2285. # The state delta.
  2286. delta: Delta = dataclasses.field(default_factory=dict)
  2287. # Events to be added to the event queue.
  2288. events: List[Event] = dataclasses.field(default_factory=list)
  2289. # Whether this is the final state update for the event.
  2290. final: bool = True
  2291. def json(self) -> str:
  2292. """Convert the state update to a JSON string.
  2293. Returns:
  2294. The state update as a JSON string.
  2295. """
  2296. return format.json_dumps(self)
  2297. class StateManager(Base, ABC):
  2298. """A class to manage many client states."""
  2299. # The state class to use.
  2300. state: Type[BaseState]
  2301. @classmethod
  2302. def create(cls, state: Type[BaseState]):
  2303. """Create a new state manager.
  2304. Args:
  2305. state: The state class to use.
  2306. Raises:
  2307. InvalidStateManagerMode: If the state manager mode is invalid.
  2308. Returns:
  2309. The state manager (either disk, memory or redis).
  2310. """
  2311. config = get_config()
  2312. if prerequisites.parse_redis_url() is not None:
  2313. config.state_manager_mode = constants.StateManagerMode.REDIS
  2314. if config.state_manager_mode == constants.StateManagerMode.MEMORY:
  2315. return StateManagerMemory(state=state)
  2316. if config.state_manager_mode == constants.StateManagerMode.DISK:
  2317. return StateManagerDisk(state=state)
  2318. if config.state_manager_mode == constants.StateManagerMode.REDIS:
  2319. redis = prerequisites.get_redis()
  2320. if redis is not None:
  2321. # make sure expiration values are obtained only from the config object on creation
  2322. return StateManagerRedis(
  2323. state=state,
  2324. redis=redis,
  2325. token_expiration=config.redis_token_expiration,
  2326. lock_expiration=config.redis_lock_expiration,
  2327. )
  2328. raise InvalidStateManagerMode(
  2329. f"Expected one of: DISK, MEMORY, REDIS, got {config.state_manager_mode}"
  2330. )
  2331. @abstractmethod
  2332. async def get_state(self, token: str) -> BaseState:
  2333. """Get the state for a token.
  2334. Args:
  2335. token: The token to get the state for.
  2336. Returns:
  2337. The state for the token.
  2338. """
  2339. pass
  2340. @abstractmethod
  2341. async def set_state(self, token: str, state: BaseState):
  2342. """Set the state for a token.
  2343. Args:
  2344. token: The token to set the state for.
  2345. state: The state to set.
  2346. """
  2347. pass
  2348. @abstractmethod
  2349. @contextlib.asynccontextmanager
  2350. async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
  2351. """Modify the state for a token while holding exclusive lock.
  2352. Args:
  2353. token: The token to modify the state for.
  2354. Yields:
  2355. The state for the token.
  2356. """
  2357. yield self.state()
  2358. class StateManagerMemory(StateManager):
  2359. """A state manager that stores states in memory."""
  2360. # The mapping of client ids to states.
  2361. states: Dict[str, BaseState] = {}
  2362. # The mutex ensures the dict of mutexes is updated exclusively
  2363. _state_manager_lock = asyncio.Lock()
  2364. # The dict of mutexes for each client
  2365. _states_locks: Dict[str, asyncio.Lock] = pydantic.PrivateAttr({})
  2366. class Config:
  2367. """The Pydantic config."""
  2368. fields = {
  2369. "_states_locks": {"exclude": True},
  2370. }
  2371. @override
  2372. async def get_state(self, token: str) -> BaseState:
  2373. """Get the state for a token.
  2374. Args:
  2375. token: The token to get the state for.
  2376. Returns:
  2377. The state for the token.
  2378. """
  2379. # Memory state manager ignores the substate suffix and always returns the top-level state.
  2380. token = _split_substate_key(token)[0]
  2381. if token not in self.states:
  2382. self.states[token] = self.state(_reflex_internal_init=True)
  2383. return self.states[token]
  2384. @override
  2385. async def set_state(self, token: str, state: BaseState):
  2386. """Set the state for a token.
  2387. Args:
  2388. token: The token to set the state for.
  2389. state: The state to set.
  2390. """
  2391. pass
  2392. @override
  2393. @contextlib.asynccontextmanager
  2394. async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
  2395. """Modify the state for a token while holding exclusive lock.
  2396. Args:
  2397. token: The token to modify the state for.
  2398. Yields:
  2399. The state for the token.
  2400. """
  2401. # Memory state manager ignores the substate suffix and always returns the top-level state.
  2402. token = _split_substate_key(token)[0]
  2403. if token not in self._states_locks:
  2404. async with self._state_manager_lock:
  2405. if token not in self._states_locks:
  2406. self._states_locks[token] = asyncio.Lock()
  2407. async with self._states_locks[token]:
  2408. state = await self.get_state(token)
  2409. yield state
  2410. await self.set_state(token, state)
  2411. def _default_token_expiration() -> int:
  2412. """Get the default token expiration time.
  2413. Returns:
  2414. The default token expiration time.
  2415. """
  2416. return get_config().redis_token_expiration
  2417. def _serialize_type(type_: Any) -> str:
  2418. """Serialize a type.
  2419. Args:
  2420. type_: The type to serialize.
  2421. Returns:
  2422. The serialized type.
  2423. """
  2424. if not inspect.isclass(type_):
  2425. return f"{type_}"
  2426. return f"{type_.__module__}.{type_.__qualname__}"
  2427. def is_serializable(value: Any) -> bool:
  2428. """Check if a value is serializable.
  2429. Args:
  2430. value: The value to check.
  2431. Returns:
  2432. Whether the value is serializable.
  2433. """
  2434. try:
  2435. return bool(pickle.dumps(value))
  2436. except Exception:
  2437. return False
  2438. def reset_disk_state_manager():
  2439. """Reset the disk state manager."""
  2440. states_directory = prerequisites.get_web_dir() / constants.Dirs.STATES
  2441. if states_directory.exists():
  2442. for path in states_directory.iterdir():
  2443. path.unlink()
  2444. class StateManagerDisk(StateManager):
  2445. """A state manager that stores states in memory."""
  2446. # The mapping of client ids to states.
  2447. states: Dict[str, BaseState] = {}
  2448. # The mutex ensures the dict of mutexes is updated exclusively
  2449. _state_manager_lock = asyncio.Lock()
  2450. # The dict of mutexes for each client
  2451. _states_locks: Dict[str, asyncio.Lock] = pydantic.PrivateAttr({})
  2452. # The token expiration time (s).
  2453. token_expiration: int = pydantic.Field(default_factory=_default_token_expiration)
  2454. class Config:
  2455. """The Pydantic config."""
  2456. fields = {
  2457. "_states_locks": {"exclude": True},
  2458. }
  2459. keep_untouched = (functools.cached_property,)
  2460. def __init__(self, state: Type[BaseState]):
  2461. """Create a new state manager.
  2462. Args:
  2463. state: The state class to use.
  2464. """
  2465. super().__init__(state=state)
  2466. path_ops.mkdir(self.states_directory)
  2467. self._purge_expired_states()
  2468. @functools.cached_property
  2469. def states_directory(self) -> Path:
  2470. """Get the states directory.
  2471. Returns:
  2472. The states directory.
  2473. """
  2474. return prerequisites.get_web_dir() / constants.Dirs.STATES
  2475. def _purge_expired_states(self):
  2476. """Purge expired states from the disk."""
  2477. import time
  2478. for path in path_ops.ls(self.states_directory):
  2479. # check path is a pickle file
  2480. if path.suffix != ".pkl":
  2481. continue
  2482. # load last edited field from file
  2483. last_edited = path.stat().st_mtime
  2484. # check if the file is older than the token expiration time
  2485. if time.time() - last_edited > self.token_expiration:
  2486. # remove the file
  2487. path.unlink()
  2488. def token_path(self, token: str) -> Path:
  2489. """Get the path for a token.
  2490. Args:
  2491. token: The token to get the path for.
  2492. Returns:
  2493. The path for the token.
  2494. """
  2495. return (
  2496. self.states_directory / f"{md5(token.encode()).hexdigest()}.pkl"
  2497. ).absolute()
  2498. async def load_state(self, token: str) -> BaseState | None:
  2499. """Load a state object based on the provided token.
  2500. Args:
  2501. token: The token used to identify the state object.
  2502. Returns:
  2503. The loaded state object or None.
  2504. """
  2505. token_path = self.token_path(token)
  2506. if token_path.exists():
  2507. try:
  2508. with token_path.open(mode="rb") as file:
  2509. return BaseState._deserialize(fp=file)
  2510. except Exception:
  2511. pass
  2512. async def populate_substates(
  2513. self, client_token: str, state: BaseState, root_state: BaseState
  2514. ):
  2515. """Populate the substates of a state object.
  2516. Args:
  2517. client_token: The client token.
  2518. state: The state object to populate.
  2519. root_state: The root state object.
  2520. """
  2521. for substate in state.get_substates():
  2522. substate_token = _substate_key(client_token, substate)
  2523. fresh_instance = await root_state.get_state(substate)
  2524. instance = await self.load_state(substate_token)
  2525. if instance is not None:
  2526. # Ensure all substates exist, even if they weren't serialized previously.
  2527. instance.substates = fresh_instance.substates
  2528. else:
  2529. instance = fresh_instance
  2530. state.substates[substate.get_name()] = instance
  2531. instance.parent_state = state
  2532. await self.populate_substates(client_token, instance, root_state)
  2533. @override
  2534. async def get_state(
  2535. self,
  2536. token: str,
  2537. ) -> BaseState:
  2538. """Get the state for a token.
  2539. Args:
  2540. token: The token to get the state for.
  2541. Returns:
  2542. The state for the token.
  2543. """
  2544. client_token = _split_substate_key(token)[0]
  2545. root_state = self.states.get(client_token)
  2546. if root_state is not None:
  2547. # Retrieved state from memory.
  2548. return root_state
  2549. # Deserialize root state from disk.
  2550. root_state = await self.load_state(_substate_key(client_token, self.state))
  2551. # Create a new root state tree with all substates instantiated.
  2552. fresh_root_state = self.state(_reflex_internal_init=True)
  2553. if root_state is None:
  2554. root_state = fresh_root_state
  2555. else:
  2556. # Ensure all substates exist, even if they were not serialized previously.
  2557. root_state.substates = fresh_root_state.substates
  2558. self.states[client_token] = root_state
  2559. await self.populate_substates(client_token, root_state, root_state)
  2560. return root_state
  2561. async def set_state_for_substate(self, client_token: str, substate: BaseState):
  2562. """Set the state for a substate.
  2563. Args:
  2564. client_token: The client token.
  2565. substate: The substate to set.
  2566. """
  2567. substate_token = _substate_key(client_token, substate)
  2568. if substate._get_was_touched():
  2569. substate._was_touched = False # Reset the touched flag after serializing.
  2570. pickle_state = substate._serialize()
  2571. if pickle_state:
  2572. if not self.states_directory.exists():
  2573. self.states_directory.mkdir(parents=True, exist_ok=True)
  2574. self.token_path(substate_token).write_bytes(pickle_state)
  2575. for substate_substate in substate.substates.values():
  2576. await self.set_state_for_substate(client_token, substate_substate)
  2577. @override
  2578. async def set_state(self, token: str, state: BaseState):
  2579. """Set the state for a token.
  2580. Args:
  2581. token: The token to set the state for.
  2582. state: The state to set.
  2583. """
  2584. client_token, substate = _split_substate_key(token)
  2585. await self.set_state_for_substate(client_token, state)
  2586. @override
  2587. @contextlib.asynccontextmanager
  2588. async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
  2589. """Modify the state for a token while holding exclusive lock.
  2590. Args:
  2591. token: The token to modify the state for.
  2592. Yields:
  2593. The state for the token.
  2594. """
  2595. # Memory state manager ignores the substate suffix and always returns the top-level state.
  2596. client_token, substate = _split_substate_key(token)
  2597. if client_token not in self._states_locks:
  2598. async with self._state_manager_lock:
  2599. if client_token not in self._states_locks:
  2600. self._states_locks[client_token] = asyncio.Lock()
  2601. async with self._states_locks[client_token]:
  2602. state = await self.get_state(token)
  2603. yield state
  2604. await self.set_state(token, state)
  2605. def _default_lock_expiration() -> int:
  2606. """Get the default lock expiration time.
  2607. Returns:
  2608. The default lock expiration time.
  2609. """
  2610. return get_config().redis_lock_expiration
  2611. class StateManagerRedis(StateManager):
  2612. """A state manager that stores states in redis."""
  2613. # The redis client to use.
  2614. redis: Redis
  2615. # The token expiration time (s).
  2616. token_expiration: int = pydantic.Field(default_factory=_default_token_expiration)
  2617. # The maximum time to hold a lock (ms).
  2618. lock_expiration: int = pydantic.Field(default_factory=_default_lock_expiration)
  2619. # The keyspace subscription string when redis is waiting for lock to be released
  2620. _redis_notify_keyspace_events: str = (
  2621. "K" # Enable keyspace notifications (target a particular key)
  2622. "g" # For generic commands (DEL, EXPIRE, etc)
  2623. "x" # For expired events
  2624. "e" # For evicted events (i.e. maxmemory exceeded)
  2625. )
  2626. # These events indicate that a lock is no longer held
  2627. _redis_keyspace_lock_release_events: Set[bytes] = {
  2628. b"del",
  2629. b"expire",
  2630. b"expired",
  2631. b"evicted",
  2632. }
  2633. async def _get_parent_state(
  2634. self, token: str, state: BaseState | None = None
  2635. ) -> BaseState | None:
  2636. """Get the parent state for the state requested in the token.
  2637. Args:
  2638. token: The token to get the state for (_substate_key).
  2639. state: The state instance to get parent state for.
  2640. Returns:
  2641. The parent state for the state requested by the token or None if there is no such parent.
  2642. """
  2643. parent_state = None
  2644. client_token, state_path = _split_substate_key(token)
  2645. parent_state_name = state_path.rpartition(".")[0]
  2646. if parent_state_name:
  2647. cached_substates = None
  2648. if state is not None:
  2649. cached_substates = [state]
  2650. # Retrieve the parent state to populate event handlers onto this substate.
  2651. parent_state = await self.get_state(
  2652. token=_substate_key(client_token, parent_state_name),
  2653. top_level=False,
  2654. get_substates=False,
  2655. cached_substates=cached_substates,
  2656. )
  2657. return parent_state
  2658. async def _populate_substates(
  2659. self,
  2660. token: str,
  2661. state: BaseState,
  2662. all_substates: bool = False,
  2663. ):
  2664. """Fetch and link substates for the given state instance.
  2665. There is no return value; the side-effect is that `state` will have `substates` populated,
  2666. and each substate will have its `parent_state` set to `state`.
  2667. Args:
  2668. token: The token to get the state for.
  2669. state: The state instance to populate substates for.
  2670. all_substates: Whether to fetch all substates or just required substates.
  2671. """
  2672. client_token, _ = _split_substate_key(token)
  2673. if all_substates:
  2674. # All substates are requested.
  2675. fetch_substates = state.get_substates()
  2676. else:
  2677. # Only _potentially_dirty_substates need to be fetched to recalc computed vars.
  2678. fetch_substates = state._potentially_dirty_substates()
  2679. tasks = {}
  2680. # Retrieve the necessary substates from redis.
  2681. for substate_cls in fetch_substates:
  2682. if substate_cls.get_name() in state.substates:
  2683. continue
  2684. substate_name = substate_cls.get_name()
  2685. tasks[substate_name] = asyncio.create_task(
  2686. self.get_state(
  2687. token=_substate_key(client_token, substate_cls),
  2688. top_level=False,
  2689. get_substates=all_substates,
  2690. parent_state=state,
  2691. )
  2692. )
  2693. for substate_name, substate_task in tasks.items():
  2694. state.substates[substate_name] = await substate_task
  2695. @override
  2696. async def get_state(
  2697. self,
  2698. token: str,
  2699. top_level: bool = True,
  2700. get_substates: bool = True,
  2701. parent_state: BaseState | None = None,
  2702. cached_substates: list[BaseState] | None = None,
  2703. ) -> BaseState:
  2704. """Get the state for a token.
  2705. Args:
  2706. token: The token to get the state for.
  2707. top_level: If true, return an instance of the top-level state (self.state).
  2708. get_substates: If true, also retrieve substates.
  2709. parent_state: If provided, use this parent_state instead of getting it from redis.
  2710. cached_substates: If provided, attach these substates to the state.
  2711. Returns:
  2712. The state for the token.
  2713. Raises:
  2714. RuntimeError: when the state_cls is not specified in the token
  2715. """
  2716. # Split the actual token from the fully qualified substate name.
  2717. _, state_path = _split_substate_key(token)
  2718. if state_path:
  2719. # Get the State class associated with the given path.
  2720. state_cls = self.state.get_class_substate(state_path)
  2721. else:
  2722. raise RuntimeError(
  2723. "StateManagerRedis requires token to be specified in the form of {token}_{state_full_name}"
  2724. )
  2725. # The deserialized or newly created (sub)state instance.
  2726. state = None
  2727. # Fetch the serialized substate from redis.
  2728. redis_state = await self.redis.get(token)
  2729. if redis_state is not None:
  2730. # Deserialize the substate.
  2731. with contextlib.suppress(StateSchemaMismatchError):
  2732. state = BaseState._deserialize(data=redis_state)
  2733. if state is None:
  2734. # Key didn't exist or schema mismatch so create a new instance for this token.
  2735. state = state_cls(
  2736. init_substates=False,
  2737. _reflex_internal_init=True,
  2738. )
  2739. # Populate parent state if missing and requested.
  2740. if parent_state is None:
  2741. parent_state = await self._get_parent_state(token, state)
  2742. # Set up Bidirectional linkage between this state and its parent.
  2743. if parent_state is not None:
  2744. parent_state.substates[state.get_name()] = state
  2745. state.parent_state = parent_state
  2746. # Avoid fetching substates multiple times.
  2747. if cached_substates:
  2748. for substate in cached_substates:
  2749. state.substates[substate.get_name()] = substate
  2750. if substate.parent_state is None:
  2751. substate.parent_state = state
  2752. # Populate substates if requested.
  2753. await self._populate_substates(token, state, all_substates=get_substates)
  2754. # To retain compatibility with previous implementation, by default, we return
  2755. # the top-level state by chasing `parent_state` pointers up the tree.
  2756. if top_level:
  2757. return state._get_root_state()
  2758. return state
  2759. @override
  2760. async def set_state(
  2761. self,
  2762. token: str,
  2763. state: BaseState,
  2764. lock_id: bytes | None = None,
  2765. ):
  2766. """Set the state for a token.
  2767. Args:
  2768. token: The token to set the state for.
  2769. state: The state to set.
  2770. lock_id: If provided, the lock_key must be set to this value to set the state.
  2771. Raises:
  2772. LockExpiredError: If lock_id is provided and the lock for the token is not held by that ID.
  2773. RuntimeError: If the state instance doesn't match the state name in the token.
  2774. """
  2775. # Check that we're holding the lock.
  2776. if (
  2777. lock_id is not None
  2778. and await self.redis.get(self._lock_key(token)) != lock_id
  2779. ):
  2780. raise LockExpiredError(
  2781. f"Lock expired for token {token} while processing. Consider increasing "
  2782. f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) "
  2783. "or use `@rx.event(background=True)` decorator for long-running tasks."
  2784. )
  2785. client_token, substate_name = _split_substate_key(token)
  2786. # If the substate name on the token doesn't match the instance name, it cannot have a parent.
  2787. if state.parent_state is not None and state.get_full_name() != substate_name:
  2788. raise RuntimeError(
  2789. f"Cannot `set_state` with mismatching token {token} and substate {state.get_full_name()}."
  2790. )
  2791. # Recursively set_state on all known substates.
  2792. tasks = []
  2793. for substate in state.substates.values():
  2794. tasks.append(
  2795. asyncio.create_task(
  2796. self.set_state(
  2797. token=_substate_key(client_token, substate),
  2798. state=substate,
  2799. lock_id=lock_id,
  2800. )
  2801. )
  2802. )
  2803. # Persist only the given state (parents or substates are excluded by BaseState.__getstate__).
  2804. if state._get_was_touched():
  2805. pickle_state = state._serialize()
  2806. if pickle_state:
  2807. await self.redis.set(
  2808. _substate_key(client_token, state),
  2809. pickle_state,
  2810. ex=self.token_expiration,
  2811. )
  2812. # Wait for substates to be persisted.
  2813. for t in tasks:
  2814. await t
  2815. @override
  2816. @contextlib.asynccontextmanager
  2817. async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
  2818. """Modify the state for a token while holding exclusive lock.
  2819. Args:
  2820. token: The token to modify the state for.
  2821. Yields:
  2822. The state for the token.
  2823. """
  2824. async with self._lock(token) as lock_id:
  2825. state = await self.get_state(token)
  2826. yield state
  2827. await self.set_state(token, state, lock_id)
  2828. @staticmethod
  2829. def _lock_key(token: str) -> bytes:
  2830. """Get the redis key for a token's lock.
  2831. Args:
  2832. token: The token to get the lock key for.
  2833. Returns:
  2834. The redis lock key for the token.
  2835. """
  2836. # All substates share the same lock domain, so ignore any substate path suffix.
  2837. client_token = _split_substate_key(token)[0]
  2838. return f"{client_token}_lock".encode()
  2839. async def _try_get_lock(self, lock_key: bytes, lock_id: bytes) -> bool | None:
  2840. """Try to get a redis lock for a token.
  2841. Args:
  2842. lock_key: The redis key for the lock.
  2843. lock_id: The ID of the lock.
  2844. Returns:
  2845. True if the lock was obtained.
  2846. """
  2847. return await self.redis.set(
  2848. lock_key,
  2849. lock_id,
  2850. px=self.lock_expiration,
  2851. nx=True, # only set if it doesn't exist
  2852. )
  2853. async def _wait_lock(self, lock_key: bytes, lock_id: bytes) -> None:
  2854. """Wait for a redis lock to be released via pubsub.
  2855. Coroutine will not return until the lock is obtained.
  2856. Args:
  2857. lock_key: The redis key for the lock.
  2858. lock_id: The ID of the lock.
  2859. Raises:
  2860. ResponseError: when the keyspace config cannot be set.
  2861. """
  2862. state_is_locked = False
  2863. lock_key_channel = f"__keyspace@0__:{lock_key.decode()}"
  2864. # Enable keyspace notifications for the lock key, so we know when it is available.
  2865. try:
  2866. await self.redis.config_set(
  2867. "notify-keyspace-events",
  2868. self._redis_notify_keyspace_events,
  2869. )
  2870. except ResponseError:
  2871. # Some redis servers only allow out-of-band configuration, so ignore errors here.
  2872. if not environment.REFLEX_IGNORE_REDIS_CONFIG_ERROR.get():
  2873. raise
  2874. async with self.redis.pubsub() as pubsub:
  2875. await pubsub.psubscribe(lock_key_channel)
  2876. while not state_is_locked:
  2877. # wait for the lock to be released
  2878. while True:
  2879. if not await self.redis.exists(lock_key):
  2880. break # key was removed, try to get the lock again
  2881. message = await pubsub.get_message(
  2882. ignore_subscribe_messages=True,
  2883. timeout=self.lock_expiration / 1000.0,
  2884. )
  2885. if message is None:
  2886. continue
  2887. if message["data"] in self._redis_keyspace_lock_release_events:
  2888. break
  2889. state_is_locked = await self._try_get_lock(lock_key, lock_id)
  2890. @contextlib.asynccontextmanager
  2891. async def _lock(self, token: str):
  2892. """Obtain a redis lock for a token.
  2893. Args:
  2894. token: The token to obtain a lock for.
  2895. Yields:
  2896. The ID of the lock (to be passed to set_state).
  2897. Raises:
  2898. LockExpiredError: If the lock has expired while processing the event.
  2899. """
  2900. lock_key = self._lock_key(token)
  2901. lock_id = uuid.uuid4().hex.encode()
  2902. if not await self._try_get_lock(lock_key, lock_id):
  2903. # Missed the fast-path to get lock, subscribe for lock delete/expire events
  2904. await self._wait_lock(lock_key, lock_id)
  2905. state_is_locked = True
  2906. try:
  2907. yield lock_id
  2908. except LockExpiredError:
  2909. state_is_locked = False
  2910. raise
  2911. finally:
  2912. if state_is_locked:
  2913. # only delete our lock
  2914. await self.redis.delete(lock_key)
  2915. async def close(self):
  2916. """Explicitly close the redis connection and connection_pool.
  2917. It is necessary in testing scenarios to close between asyncio test cases
  2918. to avoid having lingering redis connections associated with event loops
  2919. that will be closed (each test case uses its own event loop).
  2920. Note: Connections will be automatically reopened when needed.
  2921. """
  2922. await self.redis.aclose(close_connection_pool=True)
  2923. def get_state_manager() -> StateManager:
  2924. """Get the state manager for the app that is currently running.
  2925. Returns:
  2926. The state manager.
  2927. """
  2928. app = getattr(prerequisites.get_app(), constants.CompileVars.APP)
  2929. return app.state_manager
  2930. class MutableProxy(wrapt.ObjectProxy):
  2931. """A proxy for a mutable object that tracks changes."""
  2932. # Methods on wrapped objects which should mark the state as dirty.
  2933. __mark_dirty_attrs__ = set(
  2934. [
  2935. "add",
  2936. "append",
  2937. "clear",
  2938. "difference_update",
  2939. "discard",
  2940. "extend",
  2941. "insert",
  2942. "intersection_update",
  2943. "pop",
  2944. "popitem",
  2945. "remove",
  2946. "reverse",
  2947. "setdefault",
  2948. "sort",
  2949. "symmetric_difference_update",
  2950. "update",
  2951. ]
  2952. )
  2953. # Methods on wrapped objects might return mutable objects that should be tracked.
  2954. __wrap_mutable_attrs__ = set(
  2955. [
  2956. "get",
  2957. "setdefault",
  2958. ]
  2959. )
  2960. # These internal attributes on rx.Base should NOT be wrapped in a MutableProxy.
  2961. __never_wrap_base_attrs__ = set(Base.__dict__) - {"set"} | set(
  2962. pydantic.BaseModel.__dict__
  2963. )
  2964. # These types will be wrapped in MutableProxy
  2965. __mutable_types__ = (
  2966. list,
  2967. dict,
  2968. set,
  2969. Base,
  2970. DeclarativeBase,
  2971. BaseModelV2,
  2972. BaseModelV1,
  2973. )
  2974. def __init__(self, wrapped: Any, state: BaseState, field_name: str):
  2975. """Create a proxy for a mutable object that tracks changes.
  2976. Args:
  2977. wrapped: The object to proxy.
  2978. state: The state to mark dirty when the object is changed.
  2979. field_name: The name of the field on the state associated with the
  2980. wrapped object.
  2981. """
  2982. super().__init__(wrapped)
  2983. self._self_state = state
  2984. self._self_field_name = field_name
  2985. def _mark_dirty(
  2986. self,
  2987. wrapped=None,
  2988. instance=None,
  2989. args=tuple(),
  2990. kwargs=None,
  2991. ) -> Any:
  2992. """Mark the state as dirty, then call a wrapped function.
  2993. Intended for use with `FunctionWrapper` from the `wrapt` library.
  2994. Args:
  2995. wrapped: The wrapped function.
  2996. instance: The instance of the wrapped function.
  2997. args: The args for the wrapped function.
  2998. kwargs: The kwargs for the wrapped function.
  2999. Returns:
  3000. The result of the wrapped function.
  3001. """
  3002. self._self_state.dirty_vars.add(self._self_field_name)
  3003. self._self_state._mark_dirty()
  3004. if wrapped is not None:
  3005. return wrapped(*args, **(kwargs or {}))
  3006. @classmethod
  3007. def _is_mutable_type(cls, value: Any) -> bool:
  3008. """Check if a value is of a mutable type and should be wrapped.
  3009. Args:
  3010. value: The value to check.
  3011. Returns:
  3012. Whether the value is of a mutable type.
  3013. """
  3014. return isinstance(value, cls.__mutable_types__)
  3015. def _wrap_recursive(self, value: Any) -> Any:
  3016. """Wrap a value recursively if it is mutable.
  3017. Args:
  3018. value: The value to wrap.
  3019. Returns:
  3020. The wrapped value.
  3021. """
  3022. # Recursively wrap mutable types, but do not re-wrap MutableProxy instances.
  3023. if self._is_mutable_type(value) and not isinstance(value, MutableProxy):
  3024. return type(self)(
  3025. wrapped=value,
  3026. state=self._self_state,
  3027. field_name=self._self_field_name,
  3028. )
  3029. return value
  3030. def _wrap_recursive_decorator(self, wrapped, instance, args, kwargs) -> Any:
  3031. """Wrap a function that returns a possibly mutable value.
  3032. Intended for use with `FunctionWrapper` from the `wrapt` library.
  3033. Args:
  3034. wrapped: The wrapped function.
  3035. instance: The instance of the wrapped function.
  3036. args: The args for the wrapped function.
  3037. kwargs: The kwargs for the wrapped function.
  3038. Returns:
  3039. The result of the wrapped function (possibly wrapped in a MutableProxy).
  3040. """
  3041. return self._wrap_recursive(wrapped(*args, **kwargs))
  3042. def __getattr__(self, __name: str) -> Any:
  3043. """Get the attribute on the proxied object and return a proxy if mutable.
  3044. Args:
  3045. __name: The name of the attribute.
  3046. Returns:
  3047. The attribute value.
  3048. """
  3049. value = super().__getattr__(__name)
  3050. if callable(value):
  3051. if __name in self.__mark_dirty_attrs__:
  3052. # Wrap special callables, like "append", which should mark state dirty.
  3053. value = wrapt.FunctionWrapper(value, self._mark_dirty)
  3054. if __name in self.__wrap_mutable_attrs__:
  3055. # Wrap methods that may return mutable objects tied to the state.
  3056. value = wrapt.FunctionWrapper(
  3057. value,
  3058. self._wrap_recursive_decorator,
  3059. )
  3060. if (
  3061. isinstance(self.__wrapped__, Base)
  3062. and __name not in self.__never_wrap_base_attrs__
  3063. and hasattr(value, "__func__")
  3064. ):
  3065. # Wrap methods called on Base subclasses, which might do _anything_
  3066. return wrapt.FunctionWrapper(
  3067. functools.partial(value.__func__, self),
  3068. self._wrap_recursive_decorator,
  3069. )
  3070. if self._is_mutable_type(value) and __name not in (
  3071. "__wrapped__",
  3072. "_self_state",
  3073. ):
  3074. # Recursively wrap mutable attribute values retrieved through this proxy.
  3075. return self._wrap_recursive(value)
  3076. return value
  3077. def __getitem__(self, key) -> Any:
  3078. """Get the item on the proxied object and return a proxy if mutable.
  3079. Args:
  3080. key: The key of the item.
  3081. Returns:
  3082. The item value.
  3083. """
  3084. value = super().__getitem__(key)
  3085. # Recursively wrap mutable items retrieved through this proxy.
  3086. return self._wrap_recursive(value)
  3087. def __iter__(self) -> Any:
  3088. """Iterate over the proxied object and return a proxy if mutable.
  3089. Yields:
  3090. Each item value (possibly wrapped in MutableProxy).
  3091. """
  3092. for value in super().__iter__():
  3093. # Recursively wrap mutable items retrieved through this proxy.
  3094. yield self._wrap_recursive(value)
  3095. def __delattr__(self, name):
  3096. """Delete the attribute on the proxied object and mark state dirty.
  3097. Args:
  3098. name: The name of the attribute.
  3099. """
  3100. self._mark_dirty(super().__delattr__, args=(name,))
  3101. def __delitem__(self, key):
  3102. """Delete the item on the proxied object and mark state dirty.
  3103. Args:
  3104. key: The key of the item.
  3105. """
  3106. self._mark_dirty(super().__delitem__, args=(key,))
  3107. def __setitem__(self, key, value):
  3108. """Set the item on the proxied object and mark state dirty.
  3109. Args:
  3110. key: The key of the item.
  3111. value: The value of the item.
  3112. """
  3113. self._mark_dirty(super().__setitem__, args=(key, value))
  3114. def __setattr__(self, name, value):
  3115. """Set the attribute on the proxied object and mark state dirty.
  3116. If the attribute starts with "_self_", then the state is NOT marked
  3117. dirty as these are internal proxy attributes.
  3118. Args:
  3119. name: The name of the attribute.
  3120. value: The value of the attribute.
  3121. """
  3122. if name.startswith("_self_"):
  3123. # Special case attributes of the proxy itself, not applied to the wrapped object.
  3124. super().__setattr__(name, value)
  3125. return
  3126. self._mark_dirty(super().__setattr__, args=(name, value))
  3127. def __copy__(self) -> Any:
  3128. """Return a copy of the proxy.
  3129. Returns:
  3130. A copy of the wrapped object, unconnected to the proxy.
  3131. """
  3132. return copy.copy(self.__wrapped__)
  3133. def __deepcopy__(self, memo=None) -> Any:
  3134. """Return a deepcopy of the proxy.
  3135. Args:
  3136. memo: The memo dict to use for the deepcopy.
  3137. Returns:
  3138. A deepcopy of the wrapped object, unconnected to the proxy.
  3139. """
  3140. return copy.deepcopy(self.__wrapped__, memo=memo)
  3141. def __reduce_ex__(self, protocol_version):
  3142. """Get the state for redis serialization.
  3143. This method is called by cloudpickle to serialize the object.
  3144. It explicitly serializes the wrapped object, stripping off the mutable proxy.
  3145. Args:
  3146. protocol_version: The protocol version.
  3147. Returns:
  3148. Tuple of (wrapped class, empty args, class __getstate__)
  3149. """
  3150. return self.__wrapped__.__reduce_ex__(protocol_version)
  3151. @serializer
  3152. def serialize_mutable_proxy(mp: MutableProxy):
  3153. """Return the wrapped value of a MutableProxy.
  3154. Args:
  3155. mp: The MutableProxy to serialize.
  3156. Returns:
  3157. The wrapped object.
  3158. """
  3159. return mp.__wrapped__
  3160. _orig_json_JSONEncoder_default = json.JSONEncoder.default
  3161. def _json_JSONEncoder_default_wrapper(self: json.JSONEncoder, o: Any) -> Any:
  3162. """Wrap JSONEncoder.default to handle MutableProxy objects.
  3163. Args:
  3164. self: the JSONEncoder instance.
  3165. o: the object to serialize.
  3166. Returns:
  3167. A JSON-able object.
  3168. """
  3169. try:
  3170. return o.__wrapped__
  3171. except AttributeError:
  3172. pass
  3173. return _orig_json_JSONEncoder_default(self, o)
  3174. json.JSONEncoder.default = _json_JSONEncoder_default_wrapper
  3175. class ImmutableMutableProxy(MutableProxy):
  3176. """A proxy for a mutable object that tracks changes.
  3177. This wrapper comes from StateProxy, and will raise an exception if an attempt is made
  3178. to modify the wrapped object when the StateProxy is immutable.
  3179. """
  3180. def _mark_dirty(
  3181. self,
  3182. wrapped=None,
  3183. instance=None,
  3184. args=tuple(),
  3185. kwargs=None,
  3186. ) -> Any:
  3187. """Raise an exception when an attempt is made to modify the object.
  3188. Intended for use with `FunctionWrapper` from the `wrapt` library.
  3189. Args:
  3190. wrapped: The wrapped function.
  3191. instance: The instance of the wrapped function.
  3192. args: The args for the wrapped function.
  3193. kwargs: The kwargs for the wrapped function.
  3194. Returns:
  3195. The result of the wrapped function.
  3196. Raises:
  3197. ImmutableStateError: if the StateProxy is not mutable.
  3198. """
  3199. if not self._self_state._is_mutable():
  3200. raise ImmutableStateError(
  3201. "Background task StateProxy is immutable outside of a context "
  3202. "manager. Use `async with self` to modify state."
  3203. )
  3204. return super()._mark_dirty(
  3205. wrapped=wrapped, instance=instance, args=args, kwargs=kwargs
  3206. )
  3207. def code_uses_state_contexts(javascript_code: str) -> bool:
  3208. """Check if the rendered Javascript uses state contexts.
  3209. Args:
  3210. javascript_code: The Javascript code to check.
  3211. Returns:
  3212. True if the code attempts to access a member of StateContexts.
  3213. """
  3214. return bool("useContext(StateContexts" in javascript_code)
  3215. def reload_state_module(
  3216. module: str,
  3217. state: Type[BaseState] = State,
  3218. ) -> None:
  3219. """Reset rx.State subclasses to avoid conflict when reloading.
  3220. Args:
  3221. module: The module to reload.
  3222. state: Recursive argument for the state class to reload.
  3223. """
  3224. for subclass in tuple(state.class_subclasses):
  3225. reload_state_module(module=module, state=subclass)
  3226. if subclass.__module__ == module and module is not None:
  3227. state.class_subclasses.remove(subclass)
  3228. state._always_dirty_substates.discard(subclass.get_name())
  3229. state._computed_var_dependencies = defaultdict(set)
  3230. state._substate_var_dependencies = defaultdict(set)
  3231. state._init_var_dependency_dicts()
  3232. state.get_class_substate.cache_clear()