pulse-framework 0.1.40__py3-none-any.whl → 0.1.42__py3-none-any.whl

pulse/hooks/init.py

@@ -0,0 +1,460 @@
+from __future__ import annotations
+
+import ast
+import ctypes
+import functools
+import inspect
+import textwrap
+import types
+from collections.abc import Callable, Sequence
+from typing import Any, Literal, cast, override
+
+from pulse.hooks.core import HookState, hooks
+
+# Storage keyed by (code object, lineno) of the `with ps.init()` call site.
+_init_hook = hooks.create("init_storage", lambda: InitState())
+
+_CAN_USE_CPYTHON = hasattr(ctypes.pythonapi, "PyFrame_LocalsToFast")
+if _CAN_USE_CPYTHON:
+	PyFrame_LocalsToFast = ctypes.pythonapi.PyFrame_LocalsToFast
+	PyFrame_LocalsToFast.argtypes = [ctypes.py_object, ctypes.c_int]
+	PyFrame_LocalsToFast.restype = None
+
+
+def previous_frame() -> types.FrameType:
+	"""Get the previous frame (caller's frame) with assertions.
+
+	This skips the frame of this helper function and its immediate caller
+	to return the actual previous frame.
+	"""
+	current = inspect.currentframe()
+	assert current is not None, "currentframe() returned None"
+	# Skip this helper function's frame
+	caller = current.f_back
+	assert caller is not None, "f_back is None"
+	# Skip the caller's frame (e.g., __enter__) to get the actual previous frame
+	frame = caller.f_back
+	assert frame is not None, "f_back.f_back is None"
+	return frame
+
+
+class InitContext:
+	"""Context that captures locals on first render and restores thereafter."""
+
+	callsite: tuple[Any, int] | None
+	frame: types.FrameType | None
+	first_render: bool
+	pre_keys: set[str]
+	saved: dict[str, Any]
+
+	def __init__(self):
+		self.callsite = None
+		self.frame = None
+		self.first_render = False
+		self.pre_keys = set()
+		self.saved = {}
+
+	def __enter__(self):
+		self.frame = previous_frame()
+		self.pre_keys = set(self.frame.f_locals.keys())
+		# Use code object to disambiguate identical line numbers in different fns.
+		self.callsite = (self.frame.f_code, self.frame.f_lineno)
+
+		storage = _init_hook().storage
+		entry = storage.get(self.callsite)
+		if entry is None:
+			self.first_render = True
+			self.saved = {}
+		else:
+			self.first_render = False
+			self.saved = entry["vars"]
+		return self
+
+	def restore_variables(self):
+		if self.first_render:
+			return
+		frame = self.frame if self.frame is not None else previous_frame()
+		frame.f_locals.update(self.saved)
+		PyFrame_LocalsToFast(frame, 1)
+
+	def save(self, values: dict[str, Any]):
+		self.saved = values
+		assert self.callsite is not None, "callsite is None"
+		storage = _init_hook().storage
+		storage[self.callsite] = {"vars": values}
+
+	def _capture_new_locals(self) -> dict[str, Any]:
+		frame = self.frame
+		assert frame is not None, "frame is None"
+		captured = {}
+		for name, value in frame.f_locals.items():
+			if name in self.pre_keys:
+				continue
+			if value is self:
+				continue
+			captured[name] = value
+		return captured
+
+	def __exit__(
+		self,
+		exc_type: type[BaseException] | None,
+		exc_value: BaseException | None,
+		exc_tb: Any,
+	) -> Literal[False]:
+		if exc_type is None:
+			captured = self._capture_new_locals()
+			assert self.callsite is not None, "callsite  None"
+			storage = _init_hook().storage
+			storage[self.callsite] = {"vars": captured}
+		self.frame = None
+		return False
+
+
+def init() -> InitContext:
+	return InitContext()
+
+
+# ---------------------------- AST rewriting -------------------------------
+
+
+class InitCPythonRewriter(ast.NodeTransformer):
+	counter: int
+	_init_names: set[str]
+	_init_modules: set[str]
+
+	def __init__(self, init_names: set[str], init_modules: set[str]):
+		super().__init__()
+		self.counter = 0
+		self._init_names = init_names
+		self._init_modules = init_modules
+
+	@override
+	def visit_With(self, node: ast.With):
+		node = cast(ast.With, self.generic_visit(node))
+		if not node.items:
+			return node
+
+		item = node.items[0]
+		if self.is_init_call(item.context_expr):
+			ctx_name = f"_init_ctx_{self.counter}"
+			self.counter += 1
+			new_item = ast.withitem(
+				context_expr=item.context_expr,
+				optional_vars=ast.Name(id=ctx_name, ctx=ast.Store()),
+			)
+
+			restore_call = ast.Expr(
+				value=ast.Call(
+					func=ast.Attribute(
+						value=ast.Name(id=ctx_name, ctx=ast.Load()),
+						attr="restore_variables",
+						ctx=ast.Load(),
+					),
+					args=[],
+					keywords=[],
+				)
+			)
+
+			new_if = ast.If(
+				test=ast.Attribute(
+					value=ast.Name(id=ctx_name, ctx=ast.Load()),
+					attr="first_render",
+					ctx=ast.Load(),
+				),
+				body=node.body,
+				orelse=[restore_call],
+			)
+
+			return ast.With(
+				items=[new_item],
+				body=[new_if],
+				type_comment=getattr(node, "type_comment", None),
+			)
+
+		return node
+
+	def is_init_call(self, expr: ast.AST) -> bool:
+		if not isinstance(expr, ast.Call):
+			return False
+		func = expr.func
+		if isinstance(func, ast.Name) and func.id in self._init_names:
+			return True
+		if (
+			isinstance(func, ast.Attribute)
+			and isinstance(func.value, ast.Name)
+			and func.value.id in self._init_modules
+			and func.attr == "init"
+		):
+			return True
+		return False
+
+
+class InitFallbackRewriter(ast.NodeTransformer):
+	"""Rewrite using explicit rebinding (portable, no LocalsToFast)."""
+
+	counter: int
+	_init_names: set[str]
+	_init_modules: set[str]
+
+	def __init__(self, init_names: set[str], init_modules: set[str]):
+		super().__init__()
+		self.counter = 0
+		self._init_names = init_names
+		self._init_modules = init_modules
+
+	@override
+	def visit_With(self, node: ast.With):
+		node = cast(ast.With, self.generic_visit(node))
+		if not node.items:
+			return node
+
+		item = node.items[0]
+		if not self.is_init_call(item.context_expr):
+			return node
+
+		ctx_name = f"_init_ctx_{self.counter}"
+		self.counter += 1
+		new_item = ast.withitem(
+			context_expr=item.context_expr,
+			optional_vars=ast.Name(id=ctx_name, ctx=ast.Store()),
+		)
+
+		assigned = _collect_assigned_names(node.body)
+
+		save_call = ast.Expr(
+			value=ast.Call(
+				func=ast.Attribute(
+					value=ast.Name(id=ctx_name, ctx=ast.Load()),
+					attr="save",
+					ctx=ast.Load(),
+				),
+				args=[
+					ast.Dict(
+						keys=[ast.Constant(n) for n in assigned],
+						values=[ast.Name(id=n, ctx=ast.Load()) for n in assigned],
+					)
+				],
+				keywords=[],
+			)
+		)
+
+		restore_assigns: Sequence[ast.stmt] = [
+			ast.Assign(
+				targets=[ast.Name(id=name, ctx=ast.Store())],
+				value=ast.Subscript(
+					value=ast.Attribute(
+						value=ast.Name(id=ctx_name, ctx=ast.Load()),
+						attr="saved",
+						ctx=ast.Load(),
+					),
+					slice=ast.Constant(name),
+					ctx=ast.Load(),
+				),
+			)
+			for name in assigned
+		]
+
+		new_if = ast.If(
+			test=ast.Attribute(
+				value=ast.Name(id=ctx_name, ctx=ast.Load()),
+				attr="first_render",
+				ctx=ast.Load(),
+			),
+			body=node.body + [save_call],
+			orelse=list(restore_assigns),
+		)
+
+		return ast.With(
+			items=[new_item],
+			body=[new_if],
+			type_comment=getattr(node, "type_comment", None),
+		)
+
+	def is_init_call(self, expr: ast.AST) -> bool:
+		if not isinstance(expr, ast.Call):
+			return False
+		func = expr.func
+		if isinstance(func, ast.Name) and func.id in self._init_names:
+			return True
+		if (
+			isinstance(func, ast.Attribute)
+			and isinstance(func.value, ast.Name)
+			and func.value.id in self._init_modules
+			and func.attr == "init"
+		):
+			return True
+		return False
+
+
+def _collect_assigned_names(body: list[ast.stmt]) -> list[str]:
+	names: set[str] = set()
+
+	def add_target(target: ast.AST):
+		if isinstance(target, ast.Name):
+			names.add(target.id)
+		elif isinstance(target, (ast.Tuple, ast.List)):
+			for elt in target.elts:
+				add_target(elt)
+
+	for stmt in body:
+		if isinstance(stmt, ast.Assign):
+			for target in stmt.targets:
+				add_target(target)
+		elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
+			names.add(stmt.target.id)
+		elif isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+			names.add(stmt.name)
+	return list(names)
+
+
+def rewrite_init_blocks(func: Callable[..., Any]) -> Callable[..., Any]:
+	"""Rewrite `with ps.init()` blocks in the provided function, if present."""
+
+	source = _get_source(func)  # raises immediately if missing
+
+	if "init" not in source:  # quick prefilter, allow alias detection later
+		return func
+
+	tree = ast.parse(source)
+
+	init_names, init_modules = _resolve_init_bindings(func)
+
+	# Remove decorators so the re-exec'd function isn't double-wrapped.
+	for node in ast.walk(tree):
+		if (
+			isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+			and node.name == func.__name__
+		):
+			node.decorator_list = []
+
+	if not _contains_ps_init(tree, init_names, init_modules):
+		return func
+
+	if _has_disallowed_control_flow(tree, init_names, init_modules):
+		raise RuntimeError(
+			"ps.init blocks cannot contain control flow (if/for/while/try/with/match)"
+		)
+
+	rewriter: ast.NodeTransformer
+	if _CAN_USE_CPYTHON:
+		rewriter = InitCPythonRewriter(init_names, init_modules)
+	else:
+		rewriter = InitFallbackRewriter(init_names, init_modules)
+
+	tree = rewriter.visit(tree)
+	ast.fix_missing_locations(tree)
+
+	filename = inspect.getsourcefile(func) or "<rewrite>"
+	compiled = compile(tree, filename=filename, mode="exec")
+
+	global_ns = dict(func.__globals__)
+	closure_vars = inspect.getclosurevars(func)
+	global_ns.update(closure_vars.nonlocals)
+	# Ensure `ps` resolves during exec.
+	if "ps" not in global_ns:
+		try:
+			import pulse as ps
+
+			global_ns["ps"] = ps
+		except Exception:
+			pass
+	local_ns: dict[str, Any] = {}
+	exec(compiled, global_ns, local_ns)
+	rewritten = local_ns.get(func.__name__) or global_ns[func.__name__]
+	functools.update_wrapper(rewritten, func)
+	return rewritten
+
+
+def _contains_ps_init(
+	tree: ast.AST, init_names: set[str], init_modules: set[str]
+) -> bool:
+	checker = _InitCallChecker(init_names, init_modules)
+	return checker.contains_init(tree)
+
+
+def _has_disallowed_control_flow(
+	tree: ast.AST, init_names: set[str], init_modules: set[str]
+) -> bool:
+	disallowed = (ast.If, ast.For, ast.While, ast.Try, ast.With, ast.Match)
+	checker = _InitCallChecker(init_names, init_modules)
+	for node in ast.walk(tree):
+		if isinstance(node, ast.With):
+			first = node.items[0] if node.items else None
+			if first and checker.is_init_call(first.context_expr):
+				continue
+		if isinstance(node, disallowed):
+			return True
+	return False
+
+
+class _InitCallChecker:
+	init_names: set[str]
+	init_modules: set[str]
+
+	def __init__(self, init_names: set[str], init_modules: set[str]):
+		self.init_names = init_names
+		self.init_modules = init_modules
+
+	def is_init_call(self, expr: ast.AST) -> bool:
+		if not isinstance(expr, ast.Call):
+			return False
+		func = expr.func
+		if isinstance(func, ast.Name) and func.id in self.init_names:
+			return True
+		if (
+			isinstance(func, ast.Attribute)
+			and isinstance(func.value, ast.Name)
+			and func.value.id in self.init_modules
+			and func.attr == "init"
+		):
+			return True
+		return False
+
+	def contains_init(self, tree: ast.AST) -> bool:
+		for node in ast.walk(tree):
+			if self.is_init_call(node):
+				return True
+		return False
+
+
+def _get_source(func: Callable[..., Any]) -> str:
+	try:
+		return textwrap.dedent(inspect.getsource(func))
+	except OSError as exc:
+		src = getattr(func, "__source__", None)
+		if src is None:
+			raise RuntimeError(
+				f"ps.init rewrite failed: unable to read source ({exc})"
+			) from exc
+		return textwrap.dedent(src)
+
+
+def _resolve_init_bindings(func: Callable[..., Any]) -> tuple[set[str], set[str]]:
+	"""Find names/modules that resolve to pulse.init in the function scope."""
+
+	init_names: set[str] = set()
+	init_modules: set[str] = set()
+
+	closure = inspect.getclosurevars(func)
+	scopes = [func.__globals__, closure.nonlocals, closure.globals]
+
+	for scope in scopes:
+		for name, val in scope.items():
+			if val is init:
+				init_names.add(name)
+			try:
+				if getattr(val, "init", None) is init:
+					init_modules.add(name)
+			except Exception:
+				continue
+
+	return init_names, init_modules
+
+
+class InitState(HookState):
+	def __init__(self) -> None:
+		self.storage: dict[tuple[Any, int], dict[str, Any]] = {}
+
+	@override
+	def dispose(self) -> None:
+		self.storage.clear()

pulse/hooks/states.py

@@ -17,48 +17,106 @@ S9 = TypeVar("S9", bound=State)
 S10 = TypeVar("S10", bound=State)
 
 
+class StateNamespace:
+	__slots__: tuple[str, ...] = ("states", "key", "called")
+	states: tuple[State, ...]
+	key: str | None
+	called: bool
+
+	def __init__(self, key: str | None) -> None:
+		self.states = ()
+		self.key = key
+		self.called = False
+
+	def ensure_not_called(self) -> None:
+		if self.called:
+			key_msg = (
+				f" with key='{self.key}'" if self.key is not None else " without a key"
+			)
+			raise RuntimeError(
+				f"`pulse.states` can only be called once per component render{key_msg}"
+			)
+
+	def get_or_create_states(
+		self, args: tuple[State | Callable[[], State], ...]
+	) -> tuple[State, ...]:
+		if len(self.states) > 0:
+			# Reuse existing states
+			existing_states = self.states
+			# Validate that the number of arguments matches
+			if len(args) != len(existing_states):
+				key_msg = (
+					f" with key='{self.key}'"
+					if self.key is not None
+					else " without a key"
+				)
+				raise RuntimeError(
+					f"`pulse.states` called with {len(args)} argument(s) but was previously "
+					+ f"called with {len(existing_states)} argument(s){key_msg}. "
+					+ "The number of arguments must remain consistent across renders."
+				)
+			# Dispose any State instances passed directly as args that aren't being used
+			existing_set = set(existing_states)
+			for arg in args:
+				if isinstance(arg, State) and arg not in existing_set:
+					try:
+						if not arg.__disposed__:
+							arg.dispose()
+					except RuntimeError:
+						# Already disposed, ignore
+						pass
+			return existing_states
+
+		# Create new states
+		instances = tuple(_instantiate_state(arg) for arg in args)
+		self.states = instances
+		return instances
+
+	def dispose(self) -> None:
+		for state in self.states:
+			try:
+				if not state.__disposed__:
+					state.dispose()
+			except RuntimeError:
+				# Already disposed, ignore
+				pass
+		self.states = ()
+
+
 class StatesHookState(HookState):
-	__slots__: tuple[str, ...] = ("initialized", "states", "key", "_called")
-	initialized: bool
-	_called: bool
+	__slots__: tuple[str, ...] = ("namespaces",)
+	namespaces: dict[str | None, StateNamespace]
 
 	def __init__(self) -> None:
 		super().__init__()
-		self.initialized = False
-		self.states: tuple[State, ...] = ()
-		self.key: str | None = None
-		self._called = False
+		self.namespaces = {}
 
 	@override
 	def on_render_start(self, render_cycle: int) -> None:
 		super().on_render_start(render_cycle)
-		self._called = False
-
-	def replace(self, states: list[State], key: str | None) -> None:
-		self.dispose_states()
-		self.states = tuple(states)
-		self.key = key
-		self.initialized = True
-
-	def dispose_states(self) -> None:
-		for state in self.states:
-			state.dispose()
-		self.states = ()
-		self.initialized = False
-		self.key = None
+		if self.namespaces:
+			for namespace in self.namespaces.values():
+				namespace.called = False
+
+	def get_namespace(self, key: str | None) -> StateNamespace:
+		if key not in self.namespaces:
+			self.namespaces[key] = StateNamespace(key)
+		return self.namespaces[key]
+
+	def get_or_create_states(
+		self, args: tuple[State | Callable[[], State], ...], key: str | None
+	) -> tuple[State, ...]:
+		namespace = self.get_namespace(key)
+		namespace.ensure_not_called()
+		result = namespace.get_or_create_states(args)
+		namespace.called = True
+		return result
 
 	@override
 	def dispose(self) -> None:
-		self.dispose_states()
-
-	def ensure_not_called(self) -> None:
-		if self._called:
-			raise RuntimeError(
-				"`pulse.states` can only be called once per component render"
-			)
-
-	def mark_called(self) -> None:
-		self._called = True
+		for namespace in self.namespaces.values():
+			namespace.dispose()
+		self.namespaces.clear()
 
 
 def _instantiate_state(arg: State | Callable[[], State]) -> State:
@@ -219,29 +277,8 @@ def states(*args: S | Callable[[], S], key: str | None = ...) -> tuple[S, ...]:
 
 
 def states(*args: State | Callable[[], State], key: str | None = None):
-	state = _states_hook()
-	state.ensure_not_called()
-
-	if not state.initialized:
-		instances = [_instantiate_state(arg) for arg in args]
-		state.replace(instances, key)
-		state.mark_called()
-		result = state.states
-		return result[0] if len(result) == 1 else result
-
-	if key is not None and key != state.key:
-		instances = [_instantiate_state(arg) for arg in args]
-		state.replace(instances, key)
-		state.mark_called()
-		result = state.states
-		return result[0] if len(result) == 1 else result
-
-	for arg in args:
-		if isinstance(arg, State):
-			arg.dispose()
-
-	state.mark_called()
-	result = state.states
+	hook_state = _states_hook()
+	result = hook_state.get_or_create_states(args, key)
 	return result[0] if len(result) == 1 else result
 
 

pulse/messages.py

@@ -175,6 +175,6 @@ class Directives(TypedDict):
 	socketio: SocketIODirectives
 
 
-class PrerenderResult(TypedDict):
+class Prerender(TypedDict):
 	views: dict[str, ServerInitMessage | None]
 	directives: Directives