diff --git a/pynecone/app.py b/pynecone/app.py index b7729ca6b0b..7370bec0f3b 100644 --- a/pynecone/app.py +++ b/pynecone/app.py @@ -368,6 +368,7 @@ async def process(app: App, event: Event) -> StateUpdate: state = app.state_manager.get_state(event.token) state.router_data = event.router_data + state.router_data[constants.RouteVar.CLIENT_TOKEN] = event.token # Preprocess the event. pre = app.preprocess(state, event) diff --git a/pynecone/constants.py b/pynecone/constants.py index 5d8cc7e22f8..1c5016c91a1 100644 --- a/pynecone/constants.py +++ b/pynecone/constants.py @@ -200,6 +200,14 @@ class PathArgType(SimpleNamespace): LIST = str("arg_list") +class RouteVar(SimpleNamespace): + """Names of variables used in the router_data dict stored in State.""" + + CLIENT_TOKEN = "token" + PATH = "pathname" + QUERY = "query" + + class RouteRegex(SimpleNamespace): """Regex used for extracting path args in path.""" diff --git a/pynecone/state.py b/pynecone/state.py index d98d035cfad..be99263adef 100644 --- a/pynecone/state.py +++ b/pynecone/state.py @@ -261,13 +261,21 @@ def _set_default_value(cls, prop: BaseVar): field.required = False field.default = default_value + def get_token(self) -> str: + """Return the token of the client associated with this state. + + Returns: + The token of the client. + """ + return self.router_data.get(constants.RouteVar.CLIENT_TOKEN, "") + def get_current_page(self) -> str: """Obtain the path of current page from the router data. Returns: The current page. """ - return self.router_data.get("pathname", "") + return self.router_data.get(constants.RouteVar.PATH, "") def get_query_params(self) -> Dict[str, str]: """Obtain the query parameters for the queried page. @@ -277,7 +285,7 @@ def get_query_params(self) -> Dict[str, str]: Returns: The dict of query parameters. """ - return self.router_data.get("query", {}) + return self.router_data.get(constants.RouteVar.QUERY, {}) @classmethod def setup_dynamic_args(cls, args: dict[str, str]): diff --git a/tests/test_state.py b/tests/test_state.py index 0d076a3c81c..9be355f0846 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -7,6 +7,7 @@ from pynecone.event import Event from pynecone.state import State from pynecone.var import BaseVar, ComputedVar +from pynecone.constants import RouteVar class Object(Base): @@ -609,3 +610,31 @@ def test_format_event_handler(): utils.format_event_handler(GrandchildState.do_nothing) # type: ignore == "test_state.child_state.grandchild_state.do_nothing" ) + + +def test_get_token(test_state): + assert test_state.get_token() == "" + + token = "b181904c-3953-4a79-dc18-ae9518c22f05" + test_state.router_data = {RouteVar.CLIENT_TOKEN: token} + + assert test_state.get_token() == token + + +def test_get_current_page(test_state): + + assert test_state.get_current_page() == "" + + route = "mypage/subpage" + test_state.router_data = {RouteVar.PATH: route} + + assert test_state.get_current_page() == route + + +def test_get_query_params(test_state): + assert test_state.get_query_params() == {} + + params = {"p1": "a", "p2": "b"} + test_state.router_data = {RouteVar.QUERY: params} + + assert test_state.get_query_params() == params