pulse-framework 0.1.39__py3-none-any.whl → 0.1.41__py3-none-any.whl

pulse/queries/query.py

@@ -0,0 +1,270 @@
+import asyncio
+import time
+from collections.abc import Awaitable, Callable, Hashable
+from dataclasses import dataclass
+from typing import (
+	TYPE_CHECKING,
+	Any,
+	Generic,
+	Literal,
+	TypeAlias,
+	TypeVar,
+	cast,
+	override,
+)
+
+from pulse.helpers import (
+	MISSING,
+	Disposable,
+	call_flexible,
+	is_pytest,
+	later,
+	maybe_await,
+)
+from pulse.reactive import AsyncEffect, Computed, Signal
+
+if TYPE_CHECKING:
+	from pulse.queries.query_observer import QueryResult
+
+T = TypeVar("T")
+QueryKey: TypeAlias = tuple[Hashable, ...]
+QueryStatus: TypeAlias = Literal["loading", "success", "error"]
+QueryFetchStatus: TypeAlias = Literal["idle", "fetching", "paused"]
+
+
+class AsyncQueryEffect(AsyncEffect):
+	"""
+	Specialized AsyncEffect for queries that synchronously sets loading state
+	when rescheduled/run.
+	"""
+
+	query: "Query[Any]"
+
+	def __init__(
+		self,
+		fn: Callable[[], Awaitable[None]],
+		query: "Query[Any]",
+		name: str | None = None,
+		lazy: bool = False,
+		deps: list[Signal[Any] | Computed[Any]] | None = None,
+	):
+		self.query = query
+		super().__init__(fn, name=name, lazy=lazy, deps=deps)
+
+	@override
+	def run(self) -> asyncio.Task[Any]:
+		# Immediately set loading state before running the effect
+		self.query.fetch_status.write("fetching")
+		return super().run()
+
+
+@dataclass(slots=True)
+class QueryConfig(Generic[T]):
+	retries: int
+	retry_delay: float
+	initial_data: T | Callable[[], T] | None
+	gc_time: float
+	on_dispose: Callable[[Any], None] | None
+
+
+RETRY_DELAY_DEFAULT = 2.0 if not is_pytest() else 0.01
+
+
+class Query(Generic[T], Disposable):
+	"""
+	Represents a single query instance in a store.
+	Manages the async effect, data/status signals, and observer tracking.
+	"""
+
+	key: QueryKey | None
+	fn: Callable[[], Awaitable[T]]
+	cfg: QueryConfig[T]
+
+	# Reactive signals for query state
+	data: Signal[T | None]
+	error: Signal[Exception | None]
+	last_updated: Signal[float]
+	status: Signal[QueryStatus]
+	fetch_status: Signal[QueryFetchStatus]
+	retries: Signal[int]
+	retry_reason: Signal[Exception | None]
+
+	_observers: "list[QueryResult[T]]"
+	_effect: AsyncEffect | None
+	_gc_handle: asyncio.TimerHandle | None
+
+	def __init__(
+		self,
+		key: QueryKey | None,
+		fn: Callable[[], Awaitable[T]],
+		retries: int = 3,
+		retry_delay: float = RETRY_DELAY_DEFAULT,
+		initial_data: T | None = MISSING,
+		gc_time: float = 300.0,
+		on_dispose: Callable[[Any], None] | None = None,
+	):
+		self.key = key
+		self.fn = fn
+		self.cfg = QueryConfig(
+			retries=retries,
+			retry_delay=retry_delay,
+			initial_data=initial_data,
+			gc_time=gc_time,
+			on_dispose=on_dispose,
+		)
+
+		# Initialize reactive signals
+		self.data = Signal(
+			None if initial_data is MISSING else initial_data, name=f"query.data({key})"
+		)
+		self.error = Signal(None, name=f"query.error({key})")
+		self.last_updated = Signal(
+			time.time() if initial_data is not MISSING else 0.0,
+			name=f"query.last_updated({key})",
+		)
+		self.status = Signal(
+			"loading" if initial_data is MISSING else "success",
+			name=f"query.status({key})",
+		)
+		self.fetch_status = Signal("idle", name=f"query.fetch_status({key})")
+		self.retries = Signal(0, name=f"query.retries({key})")
+		self.retry_reason = Signal(None, name=f"query.retry_reason({key})")
+
+		self._observers = []
+		self._gc_handle = None
+		# Effect is created lazily on first observation
+		self._effect = None
+
+	def set_data(self, data: T):
+		self._set_success(data, manual=True)
+
+	def set_error(self, error: Exception):
+		self._set_error(error, manual=True)
+
+	def _set_success(self, data: T, manual: bool = False):
+		self.data.write(data)
+		self.last_updated.write(time.time())
+		self.error.write(None)
+		self.status.write("success")
+		if not manual:
+			self.fetch_status.write("idle")
+			self.retries.write(0)
+			self.retry_reason.write(None)
+
+	def _set_error(self, error: Exception, manual: bool = False):
+		self.error.write(error)
+		self.last_updated.write(time.time())
+		self.status.write("error")
+		if not manual:
+			self.fetch_status.write("idle")
+			# Don't reset retries on final error - preserve for debugging
+			# retry_reason is updated to the final error in _run
+
+	def _failed_retry(self, reason: Exception):
+		self.retries.write(self.retries.read() + 1)
+		self.retry_reason.write(reason)
+
+	@property
+	def effect(self) -> AsyncEffect:
+		"""Lazy property that creates the query effect on first access."""
+		if self._effect is None:
+			self._effect = AsyncQueryEffect(
+				self._run,
+				query=self,
+				name=f"query_effect({self.key})",
+				deps=[] if self.key is not None else None,
+			)
+		return self._effect
+
+	async def _run(self):
+		# Reset retries at start of run
+		self.retries.write(0)
+		self.retry_reason.write(None)
+
+		while True:
+			try:
+				result = await self.fn()
+				self._set_success(result)
+				for obs in self._observers:
+					if obs._on_success:  # pyright: ignore[reportPrivateUsage]
+						await maybe_await(call_flexible(obs._on_success, result))  # pyright: ignore[reportPrivateUsage]
+				return
+			except asyncio.CancelledError:
+				raise
+			except Exception as e:
+				current_retries = self.retries.read()
+				if current_retries < self.cfg.retries:
+					# Record failed retry attempt and retry
+					self._failed_retry(e)
+					# Wait before retrying
+					await asyncio.sleep(self.cfg.retry_delay)
+				else:
+					# All retries exhausted - update retry_reason to final error
+					self.retry_reason.write(e)
+					self._set_error(e)
+					for obs in self._observers:
+						if obs._on_error:  # pyright: ignore[reportPrivateUsage]
+							await maybe_await(call_flexible(obs._on_error, e))  # pyright: ignore[reportPrivateUsage]
+					return
+
+	async def refetch(self, cancel_refetch: bool = True) -> T:
+		"""
+		Reruns the query and returns the result.
+		If cancel_refetch is True (default), cancels any in-flight request and starts a new one.
+		If cancel_refetch is False, deduplicates requests if one is already in flight.
+		"""
+		if cancel_refetch:
+			self.effect.cancel()
+		return await self.wait()
+
+	async def wait(self) -> T:
+		await self.effect.wait()
+		return cast(T, self.data.read())
+
+	def invalidate(self, cancel_refetch: bool = False):
+		"""
+		Marks query as stale. If there are active observers, triggers a refetch.
+		"""
+		should_schedule = not self.effect.is_scheduled or cancel_refetch
+		if should_schedule and len(self._observers) > 0:
+			self.effect.schedule()
+
+	def observe(
+		self,
+		observer: "QueryResult[T]",
+	):
+		_ = self.effect  # ensure effect is created
+		self._observers.append(observer)
+		self.cancel_gc()
+		if observer._gc_time > 0:  # pyright: ignore[reportPrivateUsage]
+			self.cfg.gc_time = max(self.cfg.gc_time, observer._gc_time)  # pyright: ignore[reportPrivateUsage]
+
+	def unobserve(self, observer: "QueryResult[T]"):
+		"""Unregister an observer. Schedules GC if no observers remain."""
+		if observer in self._observers:
+			self._observers.remove(observer)
+		if len(self._observers) == 0:
+			self.schedule_gc()
+
+	def schedule_gc(self):
+		self.cancel_gc()
+		if self.cfg.gc_time > 0:
+			self._gc_handle = later(self.cfg.gc_time, self.dispose)
+		else:
+			self.dispose()
+
+	def cancel_gc(self):
+		if self._gc_handle:
+			self._gc_handle.cancel()
+			self._gc_handle = None
+
+	@override
+	def dispose(self):
+		"""
+		Cleans up the query entry, removing it from the store.
+		"""
+		if self._effect:
+			self._effect.dispose()
+
+		if self.cfg.on_dispose:
+			self.cfg.on_dispose(self)

pulse/queries/query_observer.py

@@ -0,0 +1,365 @@
+import inspect
+import time
+from collections.abc import Awaitable, Callable
+from typing import (
+	Any,
+	Generic,
+	TypeVar,
+	cast,
+	override,
+)
+
+from pulse.context import PulseContext
+from pulse.helpers import MISSING, Disposable
+from pulse.queries.common import OnErrorFn, OnSuccessFn, bind_state
+from pulse.queries.query import (
+	RETRY_DELAY_DEFAULT,
+	Query,
+	QueryFetchStatus,
+	QueryKey,
+	QueryStatus,
+)
+from pulse.reactive import Computed, Effect, Untrack
+from pulse.state import InitializableProperty, State
+
+T = TypeVar("T")
+TState = TypeVar("TState", bound=State)
+
+
+class QueryResult(Generic[T], Disposable):
+	"""
+	Thin wrapper around Query that adds callbacks, staleness tracking,
+	and observation lifecycle.
+
+	For keyed queries, uses a Computed to resolve the correct query based on the key.
+	"""
+
+	_query: Computed[Query[T]]
+	_stale_time: float
+	_gc_time: float
+	_keep_previous_data: bool
+	_on_success: Callable[[T], Awaitable[None] | None] | None
+	_on_error: Callable[[Exception], Awaitable[None] | None] | None
+	_callback_effect: Effect
+	_observe_effect: Effect
+	_data_computed: Computed[T | None]
+	_disposed_data: T | None
+
+	def __init__(
+		self,
+		query: Computed[Query[T]],
+		stale_time: float = 0.0,
+		gc_time: float = 300.0,
+		keep_previous_data: bool = False,
+		on_success: Callable[[T], Awaitable[None] | None] | None = None,
+		on_error: Callable[[Exception], Awaitable[None] | None] | None = None,
+	):
+		self._query = query
+		self._stale_time = stale_time
+		self._gc_time = gc_time
+		self._keep_previous_data = keep_previous_data
+		self._on_success = on_success
+		self._on_error = on_error
+		self._disposed_data = None
+
+		def observe_effect():
+			query = self._query()
+			with Untrack():
+				# This may create an effect, which should live independently of our observe effect
+				query.observe(self)
+
+			# If stale or loading, schedule refetch
+			if self.is_stale():
+				query.invalidate()
+
+			# Return cleanup function that captures the observer
+			return lambda: query.unobserve(self)
+
+		self._observe_effect = Effect(
+			observe_effect,
+			name=f"query_observe({self._query().key})",
+			immediate=True,
+		)
+		self._data_computed = Computed(
+			self._data_computed_fn, name=f"query_data({self._query().key})"
+		)
+
+	@property
+	def status(self) -> QueryStatus:
+		return self._query().status()
+
+	@property
+	def fetch_status(self) -> QueryFetchStatus:
+		return self._query().fetch_status()
+
+	# Forward property reads to the query's signals (with automatic reactive tracking)
+	@property
+	def is_loading(self) -> bool:
+		return self.status == "loading"
+
+	@property
+	def is_success(self) -> bool:
+		return self.status == "success"
+
+	@property
+	def is_error(self) -> bool:
+		return self.status == "error"
+
+	@property
+	def is_fetching(self) -> bool:
+		return self.fetch_status == "fetching"
+
+	@property
+	def error(self) -> Exception | None:
+		return self._query().error.read()
+
+	def _data_computed_fn(self, prev: T | None) -> T | None:
+		query = self._query()
+		if self._keep_previous_data and query.status() != "success":
+			return prev
+		return query.data()
+
+	@property
+	def data(self) -> T | None:
+		return self._data_computed()
+
+	@property
+	def has_loaded(self) -> bool:
+		return self.status in ("success", "error")
+
+	def is_stale(self) -> bool:
+		"""Check if the query data is stale based on stale_time."""
+		query = self._query()
+		return (time.time() - query.last_updated.read()) > self._stale_time
+
+	async def refetch(self, cancel_refetch: bool = True) -> T:
+		"""Refetch the query data."""
+		return await self._query().refetch(cancel_refetch=cancel_refetch)
+
+	async def wait(self) -> T:
+		return await self._query().wait()
+
+	def invalidate(self):
+		"""Mark the query as stale and refetch if there are observers."""
+		query = self._query()
+		query.invalidate()
+
+	def set_data(self, data: T):
+		"""Optimistically set data without changing loading/error state."""
+		query = self._query()
+		query.data.write(data)
+
+	@override
+	def dispose(self):
+		"""Clean up the result and its observe effect."""
+		self._observe_effect.dispose()
+
+
+class QueryProperty(Generic[T, TState], InitializableProperty):
+	"""
+	Descriptor for state-bound queries.
+
+	Usage:
+	    class S(ps.State):
+	        @ps.query()
+	        async def user(self) -> User: ...
+
+	        @user.key
+	        def _user_key(self):
+	            return ("user", self.user_id)
+	"""
+
+	name: str
+	_fetch_fn: "Callable[[TState], Awaitable[T]]"
+	_keep_alive: bool
+	_keep_previous_data: bool
+	_stale_time: float
+	_gc_time: float
+	_retries: int
+	_retry_delay: float
+	_initial_data: T | Callable[[TState], T] | None
+	_key_fn: Callable[[TState], QueryKey] | None
+	# Not using OnSuccessFn and OnErrorFn since unions of callables are not well
+	# supported in the type system. We just need to be careful to use
+	# call_flexible to invoke these functions.
+	_on_success_fn: Callable[[TState, T], Any] | None
+	_on_error_fn: Callable[[TState, Exception], Any] | None
+	_priv_result: str
+
+	def __init__(
+		self,
+		name: str,
+		fetch_fn: "Callable[[TState], Awaitable[T]]",
+		keep_previous_data: bool = False,
+		stale_time: float = 0.0,
+		gc_time: float = 300.0,
+		retries: int = 3,
+		retry_delay: float = RETRY_DELAY_DEFAULT,
+		initial: T | Callable[[TState], T] | None = MISSING,
+		key: Callable[[TState], QueryKey] | None = None,
+		on_success: OnSuccessFn[TState, T] | None = None,
+		on_error: OnErrorFn[TState] | None = None,
+	):
+		self.name = name
+		self._fetch_fn = fetch_fn
+		self._key_fn = None
+		self._on_success_fn = on_success  # pyright: ignore[reportAttributeAccessIssue]
+		self._on_error_fn = on_error  # pyright: ignore[reportAttributeAccessIssue]
+		self._keep_previous_data = keep_previous_data
+		self._stale_time = stale_time
+		self._gc_time = gc_time
+		self._retries = retries
+		self._retry_delay = retry_delay
+		self._initial_data = initial
+		self._priv_result = f"__query_{name}"
+
+	# Decorator to attach a key function
+	def key(self, fn: Callable[[TState], QueryKey]):
+		if self._key_fn is not None:
+			raise RuntimeError(
+				f"Duplicate key() decorator for query '{self.name}'. Only one is allowed."
+			)
+		self._key_fn = fn
+		return fn
+
+	# Decorator to attach a function providing initial data
+	def initial_data(self, fn: Callable[[TState], T]):
+		if self._initial_data is not MISSING:
+			raise RuntimeError(
+				f"Duplicate initial_data() decorator for query '{self.name}'. Only one is allowed."
+			)
+		self._initial_data = fn
+		return fn
+
+	# Decorator to attach an on-success handler (sync or async)
+	def on_success(self, fn: OnSuccessFn[TState, T]):
+		if self._on_success_fn is not None:
+			raise RuntimeError(
+				f"Duplicate on_success() decorator for query '{self.name}'. Only one is allowed."
+			)
+		self._on_success_fn = fn  # pyright: ignore[reportAttributeAccessIssue]
+		return fn
+
+	# Decorator to attach an on-error handler (sync or async)
+	def on_error(self, fn: OnErrorFn[TState]):
+		if self._on_error_fn is not None:
+			raise RuntimeError(
+				f"Duplicate on_error() decorator for query '{self.name}'. Only one is allowed."
+			)
+		self._on_error_fn = fn  # pyright: ignore[reportAttributeAccessIssue]
+		return fn
+
+	@override
+	def initialize(self, state: Any, name: str) -> QueryResult[T]:
+		# Return cached query instance if present
+		result: QueryResult[T] | None = getattr(state, self._priv_result, None)
+		if result:
+			# Don't re-initialize, just return the cached instance
+			return result
+
+		# Bind methods to this instance
+		fetch_fn = bind_state(state, self._fetch_fn)
+		initial_data = cast(
+			T | None,
+			(
+				self._initial_data(state)
+				if callable(self._initial_data)
+				and len(inspect.signature(self._initial_data).parameters) == 1
+				else self._initial_data
+			),
+		)
+
+		if self._key_fn:
+			# Keyed query: use session-wide QueryStore
+			query = self._resolve_keyed(state, fetch_fn, initial_data)
+		else:
+			# Unkeyed query: create private Query
+			query = self._resolve_unkeyed(fetch_fn, initial_data)
+
+		# Wrap query in QueryResult
+		result = QueryResult[T](
+			query=query,
+			stale_time=self._stale_time,
+			keep_previous_data=self._keep_previous_data,
+			gc_time=self._gc_time,
+			on_success=bind_state(state, self._on_success_fn)
+			if self._on_success_fn
+			else None,
+			on_error=bind_state(state, self._on_error_fn)
+			if self._on_error_fn
+			else None,
+		)
+
+		# Store result on the instance
+		setattr(state, self._priv_result, result)
+		return result
+
+	def _resolve_keyed(
+		self,
+		state: TState,
+		fetch_fn: Callable[[], Awaitable[T]],
+		initial_data: T | None,
+	) -> Computed[Query[T]]:
+		"""Create or get a keyed query from the session store using a Computed."""
+		assert self._key_fn is not None
+
+		key_computed = Computed(
+			bind_state(state, self._key_fn), name=f"query.key.{self.name}"
+		)
+		render = PulseContext.get().render
+		if render is None:
+			raise RuntimeError("No render session available")
+		store = render.query_store
+
+		def query() -> Query[T]:
+			key = key_computed()
+			return store.ensure(
+				key,
+				fetch_fn,
+				initial_data,
+				gc_time=self._gc_time,
+				retries=self._retries,
+				retry_delay=self._retry_delay,
+			)
+
+		return Computed(query, name=f"query.{self.name}")
+
+	def _resolve_unkeyed(
+		self,
+		fetch_fn: Callable[[], Awaitable[T]],
+		initial_data: T | None,
+	) -> Computed[Query[T]]:
+		"""Create a private unkeyed query."""
+		query = Query[T](
+			key=None,
+			fn=fetch_fn,
+			initial_data=initial_data,
+			gc_time=self._gc_time,
+			retries=self._retries,
+			retry_delay=self._retry_delay,
+		)
+		return Computed(lambda: query, name=f"query.{self.name}")
+
+	def __get__(self, obj: Any, objtype: Any = None) -> QueryResult[T]:
+		if obj is None:
+			return self  # pyright: ignore[reportReturnType]
+		return self.initialize(obj, self.name)
+
+
+class QueryResultWithInitial(QueryResult[T]):
+	@property
+	@override
+	def data(self) -> T:
+		return cast(T, super().data)
+
+	@property
+	@override
+	def has_loaded(self) -> bool:  # mirror base for completeness
+		return super().has_loaded
+
+
+class QueryPropertyWithInitial(QueryProperty[T, TState]):
+	@override
+	def __get__(self, obj: Any, objtype: Any = None) -> QueryResultWithInitial[T]:
+		# Reuse base initialization but narrow the return type for type-checkers
+		return cast(QueryResultWithInitial[T], super().__get__(obj, objtype))

pulse/queries/store.py

@@ -0,0 +1,60 @@
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
+
+from pulse.queries.query import RETRY_DELAY_DEFAULT, Query, QueryKey
+
+T = TypeVar("T")
+
+
+class QueryStore:
+	"""
+	Store for query entries. Manages creation, retrieval, and disposal of queries.
+	"""
+
+	def __init__(self):
+		self._entries: dict[QueryKey, Query[Any]] = {}
+
+	def get(self, key: QueryKey) -> Query[Any] | None:
+		return self._entries.get(key)
+
+	def ensure(
+		self,
+		key: QueryKey,
+		fetch_fn: Callable[[], Awaitable[T]],
+		initial_data: T | None = None,
+		gc_time: float = 300.0,
+		retries: int = 3,
+		retry_delay: float = RETRY_DELAY_DEFAULT,
+	) -> Query[T]:
+		# Return existing entry if present
+		if key in self._entries:
+			return self._entries[key]
+
+		def _on_dispose(e: Query[Any]) -> None:
+			if e.key in self._entries and self._entries[e.key] is e:
+				del self._entries[e.key]
+
+		entry = Query(
+			key,
+			fetch_fn,
+			initial_data=initial_data,
+			gc_time=gc_time,
+			retries=retries,
+			retry_delay=retry_delay,
+			on_dispose=_on_dispose,
+		)
+		self._entries[key] = entry
+		return entry
+
+	def remove(self, key: QueryKey):
+		entry = self._entries.get(key)
+		if entry:
+			entry.dispose()
+
+	def get_queries(
+		self, predicate: Callable[[Query[Any]], bool] | None = None
+	) -> list[Query[Any]]:
+		"""Get all queries matching the predicate."""
+		if predicate is None:
+			return list(self._entries.values())
+		return [e for e in self._entries.values() if predicate(e)]