pulse-framework 0.1.41__py3-none-any.whl → 0.1.43__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()