Queries

QueryKey: TypeAlias = tuple[Hashable, ...]

QueryStatus: TypeAlias = Literal["loading", "success", "error"]

ActionResult: TypeAlias = ActionSuccess[T] | ActionError

@dataclass(slots=True, frozen=True)
class ActionSuccess(Generic[T]):
    data: T
    status: Literal["success"] = "success"

@dataclass(slots=True, frozen=True)
class ActionError:
    error: Exception
    status: Literal["error"] = "error"

@overload
def query(
    fn: Callable[[TState], Awaitable[T]],
    *,
    stale_time: float = 0.0,
    gc_time: float | None = 300.0,
    refetch_interval: float | None = None,
    keep_previous_data: bool = False,
    retries: int = 3,
    retry_delay: float | None = None,
    initial_data_updated_at: float | datetime | None = None,
    enabled: bool = True,
    fetch_on_mount: bool = True,
    key: QueryKey | None = None,
) -> QueryProperty[T, TState]: ...

class UserState(ps.State):
    user_id: str = ""

    @ps.query
    async def user(self) -> User:
        return await api.get_user(self.user_id)

class UserState(ps.State):
    user_id: str = ""

    @ps.query
    async def user(self) -> User:
        return await api.get_user(self.user_id)

    @user.key
    def _user_key(self):
        return ("user", self.user_id)

@ps.query(key=("users", "current"))
async def current_user(self) -> User:
    return await api.get_current_user()

@my_query.key
def _key(self) -> QueryKey:
    return ("resource", self.id)

@my_query.initial_data
def _initial(self) -> T:
    return default_value

@my_query.on_success
def _on_success(self, data: T):
    self.last_fetched = time.time()

@my_query.on_error
def _on_error(self, error: Exception):
    logger.error(f"Failed: {error}")

result = await state.user.refetch()
if result.status == "success":
    print(result.data)

state.user.invalidate()

state.user.set_data(lambda prev: {**prev, "name": "Updated"})

@overload
def mutation(
    fn: Callable[Concatenate[TState, P], Awaitable[T]],
) -> MutationProperty[T, TState, P]: ...

class UserState(ps.State):
    @ps.mutation
    async def update_name(self, name: str) -> User:
        return await api.update_user(name=name)

@update_name.on_success
def _on_success(self, data: User):
    self.user.invalidate()

@update_name.on_error
def _on_error(self, error: Exception):
    logger.error(f"Update failed: {error}")

result = await state.update_name("New Name")

@overload
def infinite_query(
    fn: Callable[[TState, TParam], Awaitable[T]],
    *,
    initial_page_param: TParam,
    max_pages: int = 0,
    stale_time: float = 0.0,
    gc_time: float | None = 300.0,
    refetch_interval: float | None = None,
    keep_previous_data: bool = False,
    retries: int = 3,
    retry_delay: float | None = None,
    initial_data_updated_at: float | datetime | None = None,
    enabled: bool = True,
    fetch_on_mount: bool = True,
    key: QueryKey = ...,  # required, or use @query.key
) -> InfiniteQueryProperty[T, TParam, TState]: ...

class FeedState(ps.State):
    @ps.infinite_query(initial_page_param=None, key=("feed",))
    async def posts(self, cursor: str | None) -> list[Post]:
        return await api.get_posts(cursor=cursor)

    @posts.get_next_page_param
    def _next_cursor(self, pages: list[Page]) -> str | None:
        return pages[-1].data[-1].id if pages and pages[-1].data else None

class FeedState(ps.State):
    feed_type: str = "recent"

    @ps.infinite_query(initial_page_param=None)
    async def posts(self, cursor: str | None) -> list[Post]:
        return await api.get_posts(feed_type=self.feed_type, cursor=cursor)

    @posts.key
    def _posts_key(self):
        return ("feed", self.feed_type)  # pages reset when feed_type changes

    @posts.get_next_page_param
    def _next_cursor(self, pages: list[Page]) -> str | None:
        return pages[-1].data[-1].id if pages and pages[-1].data else None

@my_query.key
def _key(self) -> QueryKey:
    return ("feed", self.feed_type)

@my_query.get_next_page_param
def _next(self, pages: list[Page[T, TParam]]) -> TParam | None:
    return pages[-1].next_cursor if pages else None

@my_query.initial_data
def _initial(self) -> list[Page[T, TParam]]:
    return [Page(data=[], param=None)]

result = await state.posts.fetch_next_page()

class Page(NamedTuple, Generic[T, TParam]):
    data: T
    param: TParam