diff --git a/marimo/_ast/app.py b/marimo/_ast/app.py index cfe4d3a5c94..f97c4c3cdee 100644 --- a/marimo/_ast/app.py +++ b/marimo/_ast/app.py @@ -259,7 +259,6 @@ def __init__(self, **kwargs: Any) -> None: self._cell_manager = CellManager(prefix=cell_prefix) self._graph = dataflow.DirectedGraph() self._execution_context: ExecutionContext | None = None - self._runner = dataflow.Runner(self._graph) self._header: str | None = None self._unparsable_code: list[str] = [] @@ -780,17 +779,23 @@ def process_data(pd, batch_size, learning_rate): async def _run_cell_async( self, cell: Cell, kwargs: dict[str, Any] ) -> tuple[Any, _Namespace]: + from marimo._runtime.runner import by_refs + self._maybe_initialize() - output, defs = await self._runner.run_cell_async( - cell._cell.cell_id, kwargs + output, defs = await by_refs.run_cell_async( + self._graph, cell._cell.cell_id, kwargs ) return output, _Namespace(defs, owner=self) def _run_cell_sync( self, cell: Cell, kwargs: dict[str, Any] ) -> tuple[Any, _Namespace]: + from marimo._runtime.runner import by_refs + self._maybe_initialize() - output, defs = self._runner.run_cell_sync(cell._cell.cell_id, kwargs) + output, defs = by_refs.run_cell_sync( + self._graph, cell._cell.cell_id, kwargs + ) return output, _Namespace(defs, owner=self) async def _set_ui_element_value( @@ -1000,11 +1005,6 @@ def set_execution_context( ) -> None: self._app._execution_context = execution_context - @property - def runner(self) -> dataflow.Runner: - self._app._maybe_initialize() - return self._app._runner - def update_config(self, updates: dict[str, Any]) -> _AppConfig: return self.config.update(updates) diff --git a/marimo/_ast/cell.py b/marimo/_ast/cell.py index 7257aed1993..0c3fbd84882 100644 --- a/marimo/_ast/cell.py +++ b/marimo/_ast/cell.py @@ -481,8 +481,12 @@ def _is_coroutine(self) -> bool: if hasattr(self, "_is_coro_cached"): return self._is_coro_cached assert self._app is not None - self._is_coro_cached: bool = self._app.runner.is_coroutine( - self._cell.cell_id + from marimo._runtime.runner import by_refs + + # Currently expensive since `graph` triggers _maybe_initialize on the + # underlying App. + self._is_coro_cached: bool = by_refs.is_coroutine( + self._app.graph, self._cell.cell_id ) return self._is_coro_cached @@ -656,6 +660,8 @@ def add(mo, x, y): refs = {**from_setup, **refs} try: + # TODO(dmadisetti): consider recomputing since caching doesn't close + # closure over the correct set of refs, but this is also expensive. if self._is_coroutine: return self._app.run_cell_async(cell=self, kwargs=refs) else: diff --git a/marimo/_messaging/tracebacks.py b/marimo/_messaging/tracebacks.py index 557d71451f5..46c2d52e59b 100644 --- a/marimo/_messaging/tracebacks.py +++ b/marimo/_messaging/tracebacks.py @@ -50,10 +50,8 @@ def write_traceback(traceback: str) -> None: # In run mode, only forward to the frontend if show_tracebacks is on. if in_run_mode and not _show_tracebacks_enabled(): return - # Strip marimo's internal executor.py frame and highlight for the UI - trimmed = _trim_traceback(traceback) sys.stderr._write_with_mimetype( - _highlight_traceback(trimmed), + _highlight_traceback(traceback), mimetype="application/vnd.marimo+traceback", ) else: @@ -65,16 +63,15 @@ def write_traceback(traceback: str) -> None: if in_run_mode and not _show_tracebacks_enabled(): sys.stderr.write(traceback) return - trimmed = _trim_traceback(traceback) broadcast_notification( CellNotification( cell_id=ctx.cell_id, console=CellOutput( channel=CellChannel.STDERR, mimetype="application/vnd.marimo+traceback", - data=trimmed + data=traceback if code_mode - else _highlight_traceback(trimmed), + else _highlight_traceback(traceback), ), ), ctx.stream, @@ -84,28 +81,6 @@ def write_traceback(traceback: str) -> None: sys.stderr.write(traceback) -def _trim_traceback(traceback: str) -> str: - """ - Skip first DefaultExecutor.execute_cell traceback item which all traces start with. - """ - - lines = traceback.split("\n") - if ( - len(lines) > 2 - and lines[0] == "Traceback (most recent call last):" - and ( - '/marimo/_runtime/executor.py", line ' in lines[1] - or '\\marimo\\_runtime\\executor.py", line ' in lines[1] - ) - and lines[1].endswith(", in execute_cell") - ): - for i in range(2, len(lines)): - if lines[i].startswith(" File "): - return "\n".join(lines[:1] + lines[i:]) - - return traceback - - def format_exception_message(exc: BaseException) -> str: """Return an exception's message, including Python's helpful hints. diff --git a/marimo/_runtime/app/script_runner.py b/marimo/_runtime/app/script_runner.py index d6648ee62cf..a092e81c0bc 100644 --- a/marimo/_runtime/app/script_runner.py +++ b/marimo/_runtime/app/script_runner.py @@ -2,7 +2,6 @@ from __future__ import annotations import asyncio -from collections import deque from typing import TYPE_CHECKING, Any from marimo._ast.names import SETUP_CELL_NAME @@ -19,16 +18,19 @@ from marimo._runtime.exceptions import ( MarimoMissingRefError, MarimoRuntimeException, + unwrap_user_exception, ) from marimo._runtime.executor import ( - ExecutionConfig, - get_executor, + Evaluator, + resolve_executor, ) from marimo._runtime.patches import ( create_main_module, extract_docstring_from_header, patch_main_module_context, ) +from marimo._runtime.runner.result import RunResult +from marimo._runtime.runner.scheduler import SequentialScheduler from marimo._types.ids import CellId_t if TYPE_CHECKING: @@ -49,7 +51,6 @@ def __init__( self.app = app self.filename = filename self._docstring = extract_docstring_from_header(app._app._header) - self.cells_cancelled: set[CellId_t] = set() self._glbls = glbls if glbls else {} # Setup cell cannot be overridden, and it's possible that some @@ -61,24 +62,21 @@ def __init__( excluded=CellId_t(SETUP_CELL_NAME), ) - self.cells_to_run: deque[CellId_t] = deque( + cells_to_run = [ cid for cid in pruned_execution_order if app.cell_manager.cell_data_at(cid).cell is not None and not self.app.graph.is_disabled(cid) - ) - self._executor = get_executor(ExecutionConfig()) + ] - def _cancel(self, cell_id: CellId_t) -> None: - cancelled = { - cid - for cid in dataflow.transitive_closure(self.app.graph, {cell_id}) - if cid in self.cells_to_run - } - for cid in cancelled: - self.app.graph.cells[cid].set_run_result_status("cancelled") - self.cells_cancelled |= cancelled + self._scheduler = SequentialScheduler(cells_to_run, self.app.graph) + self._evaluator = Evaluator(executor=resolve_executor(), lifecycles=[]) + # _run_synchronous and _run_asynchronous are deliberate near-twins: + # the only difference is the await on the cell step. Keeping them + # as separate methods (rather than wrapping with asyncio.run + # unconditionally) preserves the no-event-loop guarantee for purely + # synchronous apps. def _run_synchronous( self, post_execute_hooks: list[Callable[[], Any]], @@ -95,30 +93,20 @@ def _run_synchronous( glbls.update(self._glbls) outputs: dict[CellId_t, Any] = {} - while self.cells_to_run: - cid = self.cells_to_run.popleft() - if cid in self.cells_cancelled: + while self._scheduler.pending(): + cid = self._scheduler.pop_cell() + if self._scheduler.cancelled(cid): continue - # Set up has already run in this case. + # Setup has already run by this point. if cid == CellId_t(SETUP_CELL_NAME): for hook in post_execute_hooks: hook() continue - cell = self.app.graph.cells[cid] with get_context().with_cell_id(cid): try: - output = self._executor.execute_cell( - cell, glbls, self.app.graph - ) - outputs[cid] = output - except MarimoRuntimeException as e: - unwrapped_exception: BaseException | None = e.__cause__ - - if isinstance(unwrapped_exception, MarimoStopError): - self._cancel(cid) - else: - raise + result = self._evaluator.evaluate_sync(cell, glbls) + self._handle_run_result(cid, result, outputs) finally: for hook in post_execute_hooks: hook() @@ -140,36 +128,55 @@ async def _run_asynchronous( glbls.update(self._glbls) outputs: dict[CellId_t, Any] = {} - - while self.cells_to_run: - cid = self.cells_to_run.popleft() - if cid in self.cells_cancelled: + while self._scheduler.pending(): + cid = self._scheduler.pop_cell() + if self._scheduler.cancelled(cid): continue - + # Setup has already run by this point. if cid == CellId_t(SETUP_CELL_NAME): for hook in post_execute_hooks: hook() continue - cell = self.app.graph.cells[cid] with get_context().with_cell_id(cid): try: - output = await self._executor.execute_cell_async( - cell, glbls, self.app.graph - ) - outputs[cid] = output - except MarimoRuntimeException as e: - unwrapped_exception: BaseException | None = e.__cause__ - - if isinstance(unwrapped_exception, MarimoStopError): - self._cancel(cid) - else: - raise + result = await self._evaluator.evaluate(cell, glbls) + self._handle_run_result(cid, result, outputs) finally: for hook in post_execute_hooks: hook() return outputs, glbls + def _handle_run_result( + self, + cid: CellId_t, + result: RunResult, + outputs: dict[CellId_t, Any], + ) -> None: + """Classify the Evaluator's RunResult; record output/cancel/raise.""" + exc = result.exception + if exc is None: + outputs[cid] = result.output + return + if not isinstance(exc, BaseException): + # Defensive check descendants, since all exceptions are expected to + # be wrapper.. + outputs[cid] = result.output + self._scheduler.cancel(cid) + return + if isinstance(exc, MarimoRuntimeException): + unwrapped = unwrap_user_exception(exc, self.app.graph) + if isinstance(unwrapped, MarimoStopError): + outputs[cid] = unwrapped.output + self._scheduler.cancel(cid) + return + if isinstance(unwrapped, MarimoMissingRefError): + name_err = unwrapped.name_error + raise ( + name_err if name_err is not None else unwrapped + ) from None + raise exc + def run(self) -> RunOutput: from marimo._runtime.context.script_context import ( initialize_script_context, @@ -213,7 +220,7 @@ def run(self) -> RunOutput: theme=get_context().marimo_config["display"]["theme"] ) - post_execute_hooks = [] + post_execute_hooks: list[Callable[[], Any]] = [] if DependencyManager.matplotlib.has(): from marimo._output.mpl import close_figures @@ -231,25 +238,9 @@ def run(self) -> RunOutput: ) return outputs, defs - # Cell runner manages the exception handling for kernel - # runner, but script runner should raise the wrapped - # exception if invoked directly. + # Raise the wrapped user exception from "None" so the stack + # trace points at the failing cell, not the runner. except MarimoRuntimeException as e: - # MarimoMissingRefError, wraps the under lying NameError - # for context, so we raise the NameError directly. - if isinstance(e.__cause__, MarimoMissingRefError): - # For type checking + sanity check - if not isinstance(e.__cause__.name_error, NameError): - raise MarimoRuntimeException( - "Unexpected error occurred while running the app. " - "Improperly wrapped MarimoMissingRefError exception. " - "Please report this issue to " - "https://github.com/marimo-team/marimo/issues" - ) from e.__cause__ - raise e.__cause__.name_error from e.__cause__ - # For all other exceptions, we raise the wrapped exception - # from "None" to indicate this is an Error propagation, and to not - # muddy the stacktrace from the failing cells themselves. raise e.__cause__ from None # type: ignore finally: if installed_script_context: diff --git a/marimo/_runtime/dataflow/__init__.py b/marimo/_runtime/dataflow/__init__.py index ebbe5cdadd2..7b9e304f7eb 100644 --- a/marimo/_runtime/dataflow/__init__.py +++ b/marimo/_runtime/dataflow/__init__.py @@ -7,7 +7,6 @@ from marimo import _loggers from marimo._ast.cell import CellImpl from marimo._runtime.dataflow.graph import DirectedGraph -from marimo._runtime.dataflow.runner import Runner from marimo._runtime.dataflow.topology import GraphTopology from marimo._runtime.dataflow.types import Edge, EdgeWithVar from marimo._types.ids import CellId_t @@ -253,7 +252,6 @@ def import_block_relatives(cid: CellId_t, children: bool) -> set[CellId_t]: "DirectedGraph", "Edge", "EdgeWithVar", - "Runner", "get_cycles", "get_import_block_relatives", "induced_subgraph", diff --git a/marimo/_runtime/dataflow/runner.py b/marimo/_runtime/dataflow/runner.py deleted file mode 100644 index 0b8bd0f737c..00000000000 --- a/marimo/_runtime/dataflow/runner.py +++ /dev/null @@ -1,152 +0,0 @@ -# Copyright 2026 Marimo. All rights reserved. -"""Runner utility for executing individual cells in a graph.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from marimo._runtime.executor import ExecutionConfig, get_executor - -if TYPE_CHECKING: - from marimo._ast.cell import CellImpl - from marimo._runtime.dataflow.graph import DirectedGraph - from marimo._types.ids import CellId_t - - -class Runner: - """Utility for running individual cells in a graph - - This class provides methods to a run a cell in the graph and obtain its - output (last expression) and the values of its defs. - - If needed, the runner will recursively compute the values of the cell's - refs by executing its ancestors. Refs can also be substituted by the - caller. - - TODO(akshayka): Add an API for caching defs across cell runs. - """ - - def __init__(self, graph: DirectedGraph) -> None: - self._graph = graph - self._executor = get_executor(ExecutionConfig()) - - @staticmethod - def _returns(cell_impl: CellImpl, glbls: dict[str, Any]) -> dict[str, Any]: - return {name: glbls[name] for name in cell_impl.defs if name in glbls} - - @staticmethod - def _substitute_refs( - cell_impl: CellImpl, - glbls: dict[str, Any], - kwargs: dict[str, Any], - ) -> None: - for argname, argvalue in kwargs.items(): - if argname in cell_impl.refs: - glbls[argname] = argvalue - else: - raise ValueError( - f"Cell got unexpected argument {argname}" - f"The allowed arguments are {cell_impl.refs}." - ) - - def _get_ancestors( - self, cell_impl: CellImpl, kwargs: dict[str, Any] - ) -> set[CellId_t]: - from marimo._runtime.dataflow import transitive_closure - - # Get the transitive closure of parents defining unsubstituted refs - graph = self._graph - substitutions = set(kwargs.keys()) - unsubstituted_refs = cell_impl.refs - substitutions - parent_ids = { - parent_id - for parent_id in graph.parents[cell_impl.cell_id] - if graph.cells[parent_id].defs.intersection(unsubstituted_refs) - } - return transitive_closure(graph, parent_ids, children=False) - - @staticmethod - def _validate_kwargs(cell_impl: CellImpl, kwargs: dict[str, Any]) -> None: - for argname in kwargs: - if argname not in cell_impl.refs: - raise ValueError( - f"Cell got unexpected argument {argname}; " - f"The allowed arguments are {cell_impl.refs}." - ) - - def is_coroutine(self, cell_id: CellId_t) -> bool: - return self._graph.cells[cell_id].is_coroutine() or any( - self._graph.cells[cid].is_coroutine() - for cid in self._get_ancestors( - self._graph.cells[cell_id], kwargs={} - ) - ) - - async def run_cell_async( - self, cell_id: CellId_t, kwargs: dict[str, Any] - ) -> tuple[Any, dict[str, Any]]: - """Run a possibly async cell and its ancestors - - Substitutes kwargs as refs for the cell, omitting ancestors that - whose refs are substituted. - """ - from marimo._runtime.dataflow import topological_sort - - graph = self._graph - cell_impl = graph.cells[cell_id] - Runner._validate_kwargs(cell_impl, kwargs) - ancestor_ids = self._get_ancestors(cell_impl, kwargs) - - glbls: dict[str, Any] = {} - for cid in topological_sort(graph, ancestor_ids): - await self._executor.execute_cell_async( - graph.cells[cid], glbls, graph - ) - - Runner._substitute_refs(cell_impl, glbls, kwargs) - output = await self._executor.execute_cell_async( - graph.cells[cell_impl.cell_id], glbls, graph - ) - defs = Runner._returns(cell_impl, glbls) - return output, defs - - def run_cell_sync( - self, cell_id: CellId_t, kwargs: dict[str, Any] - ) -> tuple[Any, dict[str, Any]]: - """Run a synchronous cell and its ancestors - - Substitutes kwargs as refs for the cell, omitting ancestors that - whose refs are substituted. - - Raises a `RuntimeError` if the cell or any of its unsubstituted - ancestors are coroutine functions. - """ - from marimo._runtime.dataflow import topological_sort - - graph = self._graph - cell_impl = graph.cells[cell_id] - if cell_impl.is_coroutine(): - raise RuntimeError( - "A coroutine function can't be run synchronously. " - "Use `run_async()` instead" - ) - - Runner._validate_kwargs(cell_impl, kwargs) - ancestor_ids = self._get_ancestors(cell_impl, kwargs) - - if any(graph.cells[cid].is_coroutine() for cid in ancestor_ids): - raise RuntimeError( - "Cell has an ancestor that is a " - "coroutine (async) cell. Use `run_async()` instead" - ) - - glbls: dict[str, Any] = {} - for cid in topological_sort(graph, ancestor_ids): - self._executor.execute_cell(graph.cells[cid], glbls, graph) - - self._substitute_refs(cell_impl, glbls, kwargs) - output = self._executor.execute_cell( - graph.cells[cell_impl.cell_id], glbls, graph - ) - defs = Runner._returns(cell_impl, glbls) - return output, defs diff --git a/marimo/_runtime/exceptions.py b/marimo/_runtime/exceptions.py index 8adf3ba5df9..fcb51f4cea0 100644 --- a/marimo/_runtime/exceptions.py +++ b/marimo/_runtime/exceptions.py @@ -1,6 +1,11 @@ # Copyright 2026 Marimo. All rights reserved. from __future__ import annotations +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from marimo._runtime.dataflow import DirectedGraph + class MarimoRuntimeException(BaseException): """Wrapper for all marimo runtime exceptions.""" @@ -19,3 +24,16 @@ def __init__(self, ref: str, name_error: NameError | None = None) -> None: super().__init__(ref) self.ref = ref self.name_error = name_error + + +def unwrap_user_exception( + exc: MarimoRuntimeException, + graph: DirectedGraph | None = None, +) -> BaseException | None: + """Extract the user exception from a `MarimoRuntimeException`.""" + cause = exc.__cause__ + if graph is not None and isinstance(cause, NameError): + name = getattr(cause, "name", None) + if name and name in graph.definitions: + return MarimoMissingRefError(name, cause) + return cause diff --git a/marimo/_runtime/executor.py b/marimo/_runtime/executor.py deleted file mode 100644 index 776e85201ef..00000000000 --- a/marimo/_runtime/executor.py +++ /dev/null @@ -1,292 +0,0 @@ -# Copyright 2026 Marimo. All rights reserved. -from __future__ import annotations - -import inspect -import re -from abc import ABC, abstractmethod -from copy import deepcopy -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -from marimo._ast.cell import CellImpl, _is_coroutine -from marimo._ast.variables import is_mangled_local -from marimo._entrypoints.registry import EntryPointRegistry -from marimo._runtime.copy import ( - CloneError, - ShallowCopy, - ZeroCopy, - shallow_copy, -) -from marimo._runtime.exceptions import ( - MarimoMissingRefError, - MarimoNameError, - MarimoRuntimeException, -) -from marimo._runtime.primitives import ( - CLONE_PRIMITIVES, - build_ref_predicate_for_primitives, - from_unclonable_module, - is_unclonable_type, -) - -if TYPE_CHECKING: - from marimo._runtime.dataflow import DirectedGraph - -_EXECUTOR_REGISTRY = EntryPointRegistry[type["Executor"]]( - "marimo.cell.executor", -) - - -def get_executor( - config: ExecutionConfig, - registry: EntryPointRegistry[type[Executor]] = _EXECUTOR_REGISTRY, -) -> Executor: - """Get a code executor based on the execution configuration.""" - executors = registry.get_all() - - base: Executor = DefaultExecutor() - if config.is_strict: - base = StrictExecutor(base) - - for executor in executors: - base = executor(base) - return base - - -@dataclass -class ExecutionConfig: - """Configuration for cell execution.""" - - is_strict: bool = False - - -def _raise_name_error( - graph: DirectedGraph | None, name_error: NameError -) -> None: - if graph is None: - raise MarimoRuntimeException from name_error - (missing_name,) = re.findall(r"'([^']*)'", str(name_error)) - # Will miss "locals" by default since not in the graph defs. - if missing_name in graph.definitions: - raise MarimoRuntimeException from MarimoMissingRefError( - missing_name, name_error - ) - raise MarimoRuntimeException from name_error - - -class Executor(ABC): - def __init__(self, base: Executor | None = None) -> None: - self.base = base - - @abstractmethod - def execute_cell( - self, - cell: CellImpl, - glbls: dict[str, Any], - graph: DirectedGraph, - ) -> Any: - pass - - @abstractmethod - async def execute_cell_async( - self, - cell: CellImpl, - glbls: dict[str, Any], - graph: DirectedGraph, - ) -> Any: - pass - - -class DefaultExecutor(Executor): - async def execute_cell_async( - self, - cell: CellImpl, - glbls: dict[str, Any], - graph: DirectedGraph | None = None, - ) -> Any: - if cell.body is None: - return None - assert cell.last_expr is not None - try: - if _is_coroutine(cell.body): - await eval(cell.body, glbls) - else: - exec(cell.body, glbls) - - if _is_coroutine(cell.last_expr): - return await eval(cell.last_expr, glbls) - else: - return eval(cell.last_expr, glbls) - except NameError as e: - _raise_name_error(graph, e) - except (BaseException, Exception) as e: - # Raising from a BaseException will fold in the stacktrace prior - # to execution - raise MarimoRuntimeException from e - - def execute_cell( - self, - cell: CellImpl, - glbls: dict[str, Any], - graph: DirectedGraph | None = None, - ) -> Any: - try: - if cell.body is None: - return None - assert cell.last_expr is not None - - exec(cell.body, glbls) - return eval(cell.last_expr, glbls) - except NameError as e: - _raise_name_error(graph, e) - except (BaseException, Exception) as e: - raise MarimoRuntimeException from e - - -class StrictExecutor(Executor): - async def execute_cell_async( - self, - cell: CellImpl, - glbls: dict[str, Any], - graph: DirectedGraph, - ) -> Any: - assert self.base is not None, "Invalid executor composition." - - # Manage globals and references, but refers to the default beyond that. - refs = graph.get_transitive_references( - cell.refs, - predicate=build_ref_predicate_for_primitives( - glbls, CLONE_PRIMITIVES - ), - ) - backup = self._sanitize_inputs(cell, refs, glbls) - try: - response = await self.base.execute_cell_async(cell, glbls, graph) - finally: - # Restore globals from backup and backfill outputs - self._update_outputs(cell, glbls, backup) - return response - - def execute_cell( - self, - cell: CellImpl, - glbls: dict[str, Any], - graph: DirectedGraph, - ) -> Any: - assert self.base is not None, "Invalid executor composition." - - refs = graph.get_transitive_references( - cell.refs, - predicate=build_ref_predicate_for_primitives( - glbls, CLONE_PRIMITIVES - ), - ) - backup = self._sanitize_inputs(cell, refs, glbls) - try: - response = self.base.execute_cell(cell, glbls, graph) - finally: - self._update_outputs(cell, glbls, backup) - return response - - def _sanitize_inputs( - self, - cell: CellImpl, - refs: set[str], - glbls: dict[str, Any], - ) -> dict[str, Any]: - # Some attributes should remain global - lcls = { - key: glbls[key] - for key in [ - "_MicropipFinder", - "_MicropipLoader", - "__builtin__", - "__doc__", - "__file__", - "__marimo__", - "__name__", - "__package__", - "__loader__", - "__spec__", - "input", - ] - if key in glbls - } - - for ref in refs: - if ref in glbls: - if ( - isinstance( - glbls[ref], - (ZeroCopy), - ) - or inspect.ismodule(glbls[ref]) - or inspect.isfunction(glbls[ref]) - or from_unclonable_module(glbls[ref]) - or is_unclonable_type(glbls[ref]) - ): - lcls[ref] = glbls[ref] - elif isinstance(glbls[ref], ShallowCopy): - lcls[ref] = shallow_copy(glbls[ref]) - else: - try: - lcls[ref] = deepcopy(glbls[ref]) - except TypeError as e: - raise CloneError( - f"Could not clone reference `{ref}` of type " - f"{getattr(glbls[ref], '__module__', '')}. " - f"{glbls[ref].__class__.__name__} " - "try wrapping the object in a `zero_copy` " - "call. If this is a common object type, consider " - "making an issue on the marimo GitHub " - "repository to never deepcopy." - ) from e - elif ref not in glbls["__builtins__"]: - if ref in cell.defs: - raise MarimoNameError( - f"name `{ref}` is referenced before definition.", ref - ) - raise MarimoMissingRefError(ref) - - # NOTE: Execution expects the globals dictionary by memory reference, - # so we need to clear it and update it with the sanitized locals, - # returning a backup of the original globals for later restoration. - # This must be performed at the end of the function to ensure valid - # state in case of failure. - backup = {**glbls} - glbls.clear() - glbls.update(lcls) - return backup - - def _update_outputs( - self, - cell: CellImpl, - glbls: dict[str, Any], - backup: dict[str, Any], - ) -> None: - # NOTE: After execution, restore global state and update outputs. - lcls = {**glbls} - glbls.clear() - glbls.update(backup) - - defs = cell.defs - for df in defs: - if df in lcls: - # Overwrite will delete the reference. - # Weak copy holds on with references. - glbls[df] = lcls[df] - # Captures the case where a variable was previously defined by the - # cell but this most recent run did not define it. The value is now - # stale and needs to be flushed. - elif df in glbls: - del glbls[df] - - # Flush all private variables from memory - for df in backup: - if is_mangled_local(df, cell.cell_id): - del glbls[df] - - # Now repopulate all private variables. - for df in lcls: - if is_mangled_local(df, cell.cell_id): - glbls[df] = lcls[df] diff --git a/marimo/_runtime/executor/__init__.py b/marimo/_runtime/executor/__init__.py new file mode 100644 index 00000000000..37e85877ea8 --- /dev/null +++ b/marimo/_runtime/executor/__init__.py @@ -0,0 +1,35 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Cell execution runtime. + +ExecutionLifecycle: Manages global information prior to execution +Executor: Runs the execution +Evaluator: Composes lifecycles and the executor. +""" + +from __future__ import annotations + +from marimo._runtime.executor.evaluator import ( + _EXECUTOR_REGISTRY, + Evaluator, + resolve_executor, +) +from marimo._runtime.executor.executor import ( + DefaultExecutor, + Executor, +) +from marimo._runtime.executor.lifecycles import ( + ExecutionLifecycle, + Skip, +) +from marimo._runtime.executor.lifecycles.strict import StrictLifecycle + +__all__ = [ + "_EXECUTOR_REGISTRY", + "DefaultExecutor", + "Evaluator", + "ExecutionLifecycle", + "Executor", + "Skip", + "StrictLifecycle", + "resolve_executor", +] diff --git a/marimo/_runtime/executor/evaluator.py b/marimo/_runtime/executor/evaluator.py new file mode 100644 index 00000000000..523e384ffa3 --- /dev/null +++ b/marimo/_runtime/executor/evaluator.py @@ -0,0 +1,243 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Evaluator composes ExecutionLifecycles around an Executor.""" + +from __future__ import annotations + +import asyncio +import contextlib +import functools +import signal +import threading +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +from marimo import _loggers +from marimo._entrypoints.registry import EntryPointRegistry +from marimo._runtime.control_flow import MarimoInterrupt +from marimo._runtime.executor.executor import DefaultExecutor, Executor +from marimo._runtime.executor.lifecycles import ExecutionLifecycle, Skip +from marimo._runtime.runner.result import RunResult +from marimo._types.globals import MutableGlobals + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from marimo._ast.cell import CellImpl + + +LOGGER = _loggers.marimo_logger() + + +class Evaluator: + """Run a cell through an `Executor`, wrapped in `ExecutionLifecycles`.""" + + def __init__( + self, + executor: Executor, + lifecycles: list[ExecutionLifecycle], + ) -> None: + self.executor = executor + self.lifecycles: list[ExecutionLifecycle] = lifecycles + + async def evaluate( + self, cell: CellImpl, glbls: MutableGlobals + ) -> RunResult: + """Compose `ExecutionLifecycle`s around an `Executor` run. + + Lifecycle `setup`s run in order; a returned `Skip` short-circuits the + body. `teardown`s run in reverse order on lifecycles whose `setup` + was reached. + """ + completed, skip, body_exc = self._setup_chain(cell, glbls) + + if body_exc is not None: + result: RunResult = RunResult(output=None, exception=body_exc) + elif skip is not None: + # Lifecycle short-circuited — pass its full RunResult through + # so `accumulated_output` and any other field survive. + result = ( + skip.result + if skip.result is not None + else RunResult(output=None, exception=None) + ) + else: + try: + value = await self.executor.execute_cell_async(cell, glbls) + result = RunResult(output=value, exception=None) + except BaseException as e: + result = RunResult(output=None, exception=e) + + return self._teardown_chain(cell, glbls, completed, result) + + def evaluate_sync( + self, cell: CellImpl, glbls: MutableGlobals + ) -> RunResult: + """`evaluate` for callers without an event loop.""" + completed, skip, body_exc = self._setup_chain(cell, glbls) + + if body_exc is not None: + result: RunResult = RunResult(output=None, exception=body_exc) + elif skip is not None: + result = ( + skip.result + if skip.result is not None + else RunResult(output=None, exception=None) + ) + else: + try: + value = self.executor.execute_cell(cell, glbls) + result = RunResult(output=value, exception=None) + except BaseException as e: + result = RunResult(output=None, exception=e) + + return self._teardown_chain(cell, glbls, completed, result) + + async def evaluate_interruptible( + self, cell: CellImpl, glbls: MutableGlobals + ) -> RunResult: + """Await `evaluate` with SIGINT capture for coroutine cells.""" + if not cell.is_coroutine(): + return await self.evaluate(cell, glbls) + future = asyncio.ensure_future(self.evaluate(cell, glbls)) + if threading.current_thread() is threading.main_thread(): + with _cancel_on_sigint(future): + return await future + return await future + + def _setup_chain( + self, cell: CellImpl, glbls: MutableGlobals + ) -> tuple[list[ExecutionLifecycle], Skip | None, BaseException | None]: + """Run each lifecycle's `setup` in order. + + Runs `setup` on each lifecycle in `self.lifecycles` until one returns a + `Skip` or one raises an exception. Returns a tuple of (completed + lifecycles, first Skip if any, first exception if any). + """ + completed: list[ExecutionLifecycle] = [] + skip: Skip | None = None + try: + for life in self.lifecycles: + decision = life.setup(cell, glbls) + completed.append(life) + if isinstance(decision, Skip): + skip = decision + break + except BaseException as e: + return completed, None, e + return completed, skip, None + + def _teardown_chain( + self, + cell: CellImpl, + glbls: MutableGlobals, + completed: list[ExecutionLifecycle], + result: RunResult, + ) -> RunResult: + """Run each lifecycle's `teardown` in order. + + Runs `teardown` on each lifecycle in `self.lifecycles` in reverse order, + on lifecycles who successfully ran `setup`. Returns a `RunResult` with + the first exception raised by a `teardown`. + """ + teardown_exc: BaseException | None = None + for life in reversed(completed): + try: + life.teardown(cell, glbls, result) + except BaseException as e: + if teardown_exc is not None: + LOGGER.error( + "teardown exception overridden by later teardown: %s", + teardown_exc, + ) + teardown_exc = e + + if teardown_exc is not None: + if result.exception is not None: + LOGGER.warning( + "body exception suppressed by teardown raise: %s", + result.exception, + ) + return replace(result, exception=teardown_exc) + return result + + +# Public entry-point registry for plugin-loaded Executors. Registered +# values are **factories** (`Callable[[], Executor]`); the kernel +# calls the factory once to get an instance and hands it to an +# `Evaluator`. +_EXECUTOR_REGISTRY: EntryPointRegistry[Callable[[], Executor]] = ( + EntryPointRegistry("marimo.cell.executor") +) + + +def resolve_executor() -> Executor: + """Return the registered executor factory's product, or `DefaultExecutor`. + + NB. Only one factory is loaded, with others logged for visibility. + """ + names = _EXECUTOR_REGISTRY.names() + if not names: + return DefaultExecutor() + name, *additional = names + if additional: + LOGGER.warning( + "multiple `marimo.cell.executor` factories registered; " + "using %r and ignoring %d other(s)", + name, + len(additional), + ) + try: + return _EXECUTOR_REGISTRY.get(name)() + except Exception as e: + LOGGER.warning( + "marimo.cell.executor factory %r failed to construct: %s; " + "falling back to `DefaultExecutor`.", + name, + e, + ) + return DefaultExecutor() + + +# Adapted from +# https://github.com/ipython/ipykernel/blob/eddd3e666a82ebec287168b0da7cfa03639a3772/ipykernel/ipkernel.py#L312 +@contextlib.contextmanager +def _cancel_on_sigint(future: asyncio.Future[Any]) -> Iterator[None]: + """Cancel `future` if a SIGINT arrives during evaluation.""" + sigint_future: asyncio.Future[int] = asyncio.Future() + + def cancel_unless_done(f: asyncio.Future[Any], _: Any) -> None: + if f.cancelled() or f.done(): + return + f.cancel() + + sigint_future.add_done_callback( + functools.partial(cancel_unless_done, future) + ) + future.add_done_callback( + functools.partial(cancel_unless_done, sigint_future) + ) + + # Capture the previously-installed SIGINT handler *before* we install + # ours so `handle_sigint` can invoke it for its side effects + # (kernel broadcast, duckdb interrupt). For async cells the actual + # halt comes from cancelling the future, not from a raised + # `MarimoInterrupt` — so we swallow that here. + prior_sigint = signal.getsignal(signal.SIGINT) + + def handle_sigint(signum: int, frame: Any) -> None: + if sigint_future.cancelled() or sigint_future.done(): + return + sigint_future.set_result(1) + if callable(prior_sigint): + try: + prior_sigint(signum, frame) + except MarimoInterrupt: + # The kernel's handler raises MarimoInterrupt for sync + # halt; we cancel the future instead. + pass + + save_sigint = signal.signal(signal.SIGINT, handle_sigint) + try: + yield + finally: + signal.signal(signal.SIGINT, save_sigint) diff --git a/marimo/_runtime/executor/executor.py b/marimo/_runtime/executor/executor.py new file mode 100644 index 00000000000..a75285feeec --- /dev/null +++ b/marimo/_runtime/executor/executor.py @@ -0,0 +1,84 @@ +# Copyright 2026 Marimo. All rights reserved. +"""An Executor executes a single cell's body.""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Any, Protocol + +from marimo._ast.cell import _is_coroutine +from marimo._runtime.exceptions import MarimoRuntimeException +from marimo._types.globals import MutableGlobals + +if TYPE_CHECKING: + from marimo._ast.cell import CellImpl + + +def _strip_frame(e: BaseException, count: int = 1) -> None: + """Drop the top `count` frames from `e.__traceback__`. + + Stops early if the traceback runs out — never strips the last + frame, so we don't lose the only frame we have. + """ + tb = e.__traceback__ + for _ in range(count): + if tb is None or tb.tb_next is None: + break + tb = tb.tb_next + e.__traceback__ = tb + + +class Executor(Protocol): + """Body strategy: how to run a cell's body.""" + + name: str + + def execute_cell(self, cell: CellImpl, glbls: MutableGlobals) -> Any: ... + + async def execute_cell_async( + self, cell: CellImpl, glbls: MutableGlobals + ) -> Any: ... + + +class DefaultExecutor: + name = "default" + + def execute_cell(self, cell: CellImpl, glbls: MutableGlobals) -> Any: + if cell.body is None: + return None + assert cell.last_expr is not None + if _is_coroutine(cell.body) or _is_coroutine(cell.last_expr): + raise RuntimeError( + "A coroutine cell cannot be run synchronously. Use " + "execute_cell_async() instead." + ) + try: + exec(cell.body, glbls) + return eval(cell.last_expr, glbls) + except asyncio.CancelledError: + # Cancellation is control flow, not user error — surface bare. + raise + except BaseException as e: + # Strip our own frame so user-facing tracebacks start at user code. + _strip_frame(e) + raise MarimoRuntimeException from e + + async def execute_cell_async( + self, cell: CellImpl, glbls: MutableGlobals + ) -> Any: + if cell.body is None: + return None + assert cell.last_expr is not None + try: + if _is_coroutine(cell.body): + await eval(cell.body, glbls) + else: + exec(cell.body, glbls) + if _is_coroutine(cell.last_expr): + return await eval(cell.last_expr, glbls) + return eval(cell.last_expr, glbls) + except asyncio.CancelledError: + raise + except BaseException as e: + _strip_frame(e) + raise MarimoRuntimeException from e diff --git a/marimo/_runtime/executor/lifecycles/__init__.py b/marimo/_runtime/executor/lifecycles/__init__.py new file mode 100644 index 00000000000..eaf5d50f977 --- /dev/null +++ b/marimo/_runtime/executor/lifecycles/__init__.py @@ -0,0 +1,61 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Per-cell execution lifecycles setup/teardown wraps around cell execution.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +from marimo._runtime.runner.result import RunResult +from marimo._types.globals import MutableGlobals + +if TYPE_CHECKING: + from marimo._ast.cell import CellImpl + + +@dataclass +class Skip: + """Returned from `ExecutionLifecycle.setup` to short-circuit the body. + + `result` is the cell's `RunResult`; lifecycles can use this to + inject a cache hit or a pre-failed result without running the body. + `None` means the lifecycle wants to skip but has no associated run + (output stays `None`, no exception). + """ + + result: RunResult | None = None + + +class ExecutionLifecycle(Protocol): + """Per-cell setup/teardown wrap for cell execution. + + `setup`s fire in composition order; `teardown`s fire in reverse + order over the lifecycles whose `setup` was reached. Per-cell state + is stashed on the instance. + """ + + name: str + + def setup(self, cell: CellImpl, glbls: MutableGlobals) -> Skip | None: + """Run before the cell body. + + May mutate `glbls`. Return `Skip` to short-circuit the body (and + any later lifecycle in the chain); return `None` to continue. + Raising is treated like a body exception — the chain unwinds via + `teardown` on the lifecycles whose `setup` was reached. + """ + ... + + def teardown( + self, + cell: CellImpl, + glbls: MutableGlobals, + run_result: RunResult, + ) -> None: + """Run after the body (or after a `Skip`), in reverse composition order. + + May mutate `glbls`. Receives the final `RunResult` from the body or + `Skip`. Raising from `teardown` replaces `run_result.exception`; the + original body exception (if any) is logged and suppressed. + """ + ... diff --git a/marimo/_runtime/executor/lifecycles/strict.py b/marimo/_runtime/executor/lifecycles/strict.py new file mode 100644 index 00000000000..383fe9ce3f3 --- /dev/null +++ b/marimo/_runtime/executor/lifecycles/strict.py @@ -0,0 +1,174 @@ +# Copyright 2026 Marimo. All rights reserved. +"""StrictLifecycle provides globals sanitization around the body.""" + +from __future__ import annotations + +import inspect +from copy import deepcopy +from typing import TYPE_CHECKING, Any + +from marimo._ast.variables import is_mangled_local, unmangle_local +from marimo._runtime.copy import ( + CloneError, + ShallowCopy, + ZeroCopy, + shallow_copy, +) +from marimo._runtime.executor.lifecycles import Skip +from marimo._runtime.primitives import ( + CLONE_PRIMITIVES, + build_ref_predicate_for_primitives, + from_unclonable_module, + is_unclonable_type, +) +from marimo._runtime.runner.result import RunResult +from marimo._types.globals import MutableGlobals + +if TYPE_CHECKING: + from marimo._ast.cell import CellImpl + from marimo._messaging.errors import MarimoStrictExecutionError + from marimo._runtime.dataflow import DirectedGraph + from marimo._types.ids import CellId_t + + +# Attributes that should remain visible inside a strict-mode cell body +# even when the rest of the globals dict is replaced by the sanitized +# transitive references. +_PRESERVED_GLOBALS: frozenset[str] = frozenset( + { + "_MicropipFinder", + "_MicropipLoader", + "__builtin__", + "__doc__", + "__file__", + "__marimo__", + "__name__", + "__package__", + "__loader__", + "__spec__", + "input", + } +) + + +class StrictLifecycle: + """Sanitize globals before exec; restore them in teardown.""" + + name = "strict" + + def __init__(self, graph: DirectedGraph) -> None: + self._graph = graph + # Per-cell setup→teardown backup. Keyed by cell_id. + self._backups: dict[CellId_t, dict[str, Any]] = {} + + def setup(self, cell: CellImpl, glbls: MutableGlobals) -> Skip | None: + refs = self._graph.get_transitive_references( + cell.refs, + predicate=build_ref_predicate_for_primitives( + glbls, CLONE_PRIMITIVES + ), + ) + + lcls = {key: glbls[key] for key in _PRESERVED_GLOBALS if key in glbls} + + for ref in refs: + if ref in glbls: + lcls[ref] = self._sanitize_ref(ref, glbls[ref]) + elif ref not in glbls["__builtins__"]: + err = self._build_strict_error(cell, ref) + return Skip(result=RunResult(output=err, exception=err)) + + # Execution expects the globals dictionary by memory reference, + # so clear it and update with the sanitized locals, stashing a + # backup for teardown. + backup = {**glbls} + glbls.clear() + glbls.update(lcls) + self._backups[cell.cell_id] = backup + return None + + def _build_strict_error( + self, cell: CellImpl, ref: str + ) -> MarimoStrictExecutionError: + """Produce the user-facing error for an unresolved ref in setup.""" + from marimo._messaging.errors import MarimoStrictExecutionError + + if ref in cell.defs: + return MarimoStrictExecutionError( + f"name `{ref}` is referenced before definition.", ref, None + ) + blamed_cell: CellId_t | None = None + try: + (blamed_cell, *_) = self._graph.get_defining_cells(ref) + except (KeyError, ValueError): + ref, var_cell_id = unmangle_local(ref) + if var_cell_id: + blamed_cell = var_cell_id + return MarimoStrictExecutionError( + f"marimo was unable to resolve a reference to `{ref}` in cell : ", + ref, + blamed_cell, + ) + + def _sanitize_ref(self, name: str, value: Any) -> Any: + if ( + isinstance(value, ZeroCopy) + or inspect.ismodule(value) + or inspect.isfunction(value) + or from_unclonable_module(value) + or is_unclonable_type(value) + ): + return value + if isinstance(value, ShallowCopy): + return shallow_copy(value) + try: + return deepcopy(value) + except TypeError as e: + raise CloneError( + f"Could not clone reference `{name}` of type " + f"{getattr(value, '__module__', '')}. " + f"{value.__class__.__name__} " + "try wrapping the object in a `zero_copy` " + "call. If this is a common object type, consider " + "making an issue on the marimo GitHub " + "repository to never deepcopy." + ) from e + + def teardown( + self, + cell: CellImpl, + glbls: MutableGlobals, + run_result: RunResult, # noqa: ARG002 + ) -> None: + backup = self._backups.pop(cell.cell_id, None) + if backup is None: + # Setup didn't complete for this cell (raised before stashing + # the backup, or a Skip earlier in the chain meant setup + # never ran). Nothing to restore. + return + + # Restore the pre-execution globals, then re-apply the cell's + # new defs over top. + lcls = {**glbls} + glbls.clear() + glbls.update(backup) + + defs = cell.defs + for df in defs: + if df in lcls: + glbls[df] = lcls[df] + elif df in glbls: + # Previously defined by this cell, not redefined this + # run — stale, flush it. + del glbls[df] + + # Flush all private variables for this cell from the restored + # backup. + for df in backup: + if is_mangled_local(df, cell.cell_id): + del glbls[df] + + # Repopulate this cell's private variables. + for df in lcls: + if is_mangled_local(df, cell.cell_id): + glbls[df] = lcls[df] diff --git a/marimo/_runtime/kernel_request_handlers.py b/marimo/_runtime/kernel_request_handlers.py index b2b76aab393..78740bc8e1b 100644 --- a/marimo/_runtime/kernel_request_handlers.py +++ b/marimo/_runtime/kernel_request_handlers.py @@ -109,8 +109,8 @@ async def _handle_execute_scratchpad( ): await self._kernel.run_scratchpad(request.code) finally: - # Always emit completion so a waiting ``ScratchCellListener`` - # doesn't block forever if ``run_scratchpad`` raises. + # Always emit completion so a waiting `ScratchCellListener` + # doesn't block forever if `run_scratchpad` raises. broadcast_notification( CompletedRunNotification(run_id=request.run_id) ) diff --git a/marimo/_runtime/runner/by_refs.py b/marimo/_runtime/runner/by_refs.py new file mode 100644 index 00000000000..4bd9ca02f07 --- /dev/null +++ b/marimo/_runtime/runner/by_refs.py @@ -0,0 +1,190 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Lightweight runner functions for use in direct cell evaluation or testing. + +Walks the cell's ancestor closure (minus any ancestor whose defs the +caller substituted via refs), runs them with a fresh globals dict, +then runs the target cell. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from marimo._runtime.control_flow import MarimoStopError +from marimo._runtime.exceptions import MarimoRuntimeException +from marimo._runtime.executor import ( + DefaultExecutor, + Evaluator, +) +from marimo._types.globals import Globals, MutableGlobals + +if TYPE_CHECKING: + from marimo._ast.cell import CellImpl + from marimo._runtime.dataflow.topology import GraphTopology + from marimo._runtime.runner.result import RunResult + from marimo._types.ids import CellId_t + + +def _new_evaluator() -> Evaluator: + """A fresh relaxed-mode Evaluator (no lifecycles).""" + return Evaluator(executor=DefaultExecutor(), lifecycles=[]) + + +def _returns(cell_impl: CellImpl, glbls: Globals) -> dict[str, Any]: + return {name: glbls[name] for name in cell_impl.defs if name in glbls} + + +def _substitute_refs( + cell_impl: CellImpl, + glbls: MutableGlobals, + refs: dict[str, Any], +) -> None: + for argname, argvalue in refs.items(): + if argname in cell_impl.refs: + glbls[argname] = argvalue + else: + raise ValueError( + f"Cell got unexpected argument {argname}" + f"The allowed arguments are {cell_impl.refs}." + ) + + +def _validate_refs(cell_impl: CellImpl, refs: dict[str, Any]) -> None: + for argname in refs: + if argname not in cell_impl.refs: + raise ValueError( + f"Cell got unexpected argument {argname}; " + f"The allowed arguments are {cell_impl.refs}." + ) + + +def _get_ancestors( + graph: GraphTopology, + cell_impl: CellImpl, + refs: dict[str, Any], +) -> set[CellId_t]: + from marimo._runtime.dataflow import transitive_closure + + substitutions = set(refs.keys()) + unsubstituted_refs = cell_impl.refs - substitutions + parent_ids = { + parent_id + for parent_id in graph.parents[cell_impl.cell_id] + if graph.cells[parent_id].defs.intersection(unsubstituted_refs) + } + return transitive_closure(graph, parent_ids, children=False) + + +def _classify(result: RunResult) -> MarimoStopError | None: + """Inspect a RunResult; raise on real errors, return the stop on mo.stop.""" + exc = result.exception + if exc is None: + return None + if isinstance(exc, MarimoStopError): + # Defensive: any caller bypassing `MarimoRuntimeException` + # wrapping (e.g. a custom Executor that raises directly) still + # gets stop-control-flow handling. + return exc + if isinstance(exc, MarimoRuntimeException) and isinstance( + exc.__cause__, MarimoStopError + ): + return exc.__cause__ + if isinstance(exc, BaseException): + raise exc + return None + + +def is_coroutine( + graph: GraphTopology, + cell_id: CellId_t, + refs: dict[str, Any] | None = None, +) -> bool: + """True if the cell or any of its unsubstituted ancestors is async. + + NB. Currently expensive due to calls on graph. + """ + return graph.cells[cell_id].is_coroutine() or any( + graph.cells[cid].is_coroutine() + for cid in _get_ancestors(graph, graph.cells[cell_id], refs or {}) + ) + + +async def run_cell_async( + graph: GraphTopology, + cell_id: CellId_t, + refs: dict[str, Any], +) -> tuple[Any, MutableGlobals]: + """Run a possibly async cell and its ancestors. + + Substitutes refs for the cell, omitting ancestors whose + refs are substituted. + """ + from marimo._runtime.dataflow import topological_sort + + cell_impl = graph.cells[cell_id] + _validate_refs(cell_impl, refs) + ancestor_ids = _get_ancestors(graph, cell_impl, refs) + + evaluator = _new_evaluator() + glbls: MutableGlobals = {} + for cid in topological_sort(graph, ancestor_ids): + stop = _classify(await evaluator.evaluate(graph.cells[cid], glbls)) + if stop is not None: + return stop.output, _returns(cell_impl, glbls) + + _substitute_refs(cell_impl, glbls, refs) + target_result = await evaluator.evaluate( + graph.cells[cell_impl.cell_id], glbls + ) + stop = _classify(target_result) + if stop is not None: + return stop.output, _returns(cell_impl, glbls) + return target_result.output, _returns(cell_impl, glbls) + + +def run_cell_sync( + graph: GraphTopology, + cell_id: CellId_t, + refs: dict[str, Any], +) -> tuple[Any, MutableGlobals]: + """Run a synchronous cell and its ancestors. + + Substitutes refs for the cell, omitting ancestors whose + refs are substituted. + + Raises `RuntimeError` if the cell or any of its unsubstituted + ancestors are coroutine functions. + """ + from marimo._runtime.dataflow import topological_sort + + cell_impl = graph.cells[cell_id] + if cell_impl.is_coroutine(): + raise RuntimeError( + "A coroutine function can't be run synchronously. " + "Use `run_async()` instead" + ) + + _validate_refs(cell_impl, refs) + ancestor_ids = _get_ancestors(graph, cell_impl, refs) + + if any(graph.cells[cid].is_coroutine() for cid in ancestor_ids): + raise RuntimeError( + "Cell has an ancestor that is a " + "coroutine (async) cell. Use `run_async()` instead" + ) + + evaluator = _new_evaluator() + glbls: MutableGlobals = {} + for cid in topological_sort(graph, ancestor_ids): + stop = _classify(evaluator.evaluate_sync(graph.cells[cid], glbls)) + if stop is not None: + return stop.output, _returns(cell_impl, glbls) + + _substitute_refs(cell_impl, glbls, refs) + target_result = evaluator.evaluate_sync( + graph.cells[cell_impl.cell_id], glbls + ) + stop = _classify(target_result) + if stop is not None: + return stop.output, _returns(cell_impl, glbls) + return target_result.output, _returns(cell_impl, glbls) diff --git a/marimo/_runtime/runner/cell_runner.py b/marimo/_runtime/runner/cell_runner.py index c878a890462..5cfec2a87d9 100644 --- a/marimo/_runtime/runner/cell_runner.py +++ b/marimo/_runtime/runner/cell_runner.py @@ -2,14 +2,8 @@ from __future__ import annotations import asyncio -import contextlib -import functools import io -import signal -import threading import traceback -from collections import deque -from dataclasses import dataclass from pathlib import Path from types import TracebackType from typing import TYPE_CHECKING, Any @@ -22,7 +16,6 @@ from marimo._messaging.errors import ( MarimoExceptionRaisedError, MarimoSQLError, - MarimoStrictExecutionError, UnknownError, ) from marimo._messaging.tracebacks import write_traceback @@ -31,12 +24,14 @@ from marimo._runtime.control_flow import MarimoInterrupt, MarimoStopError from marimo._runtime.exceptions import ( MarimoMissingRefError, - MarimoNameError, MarimoRuntimeException, + unwrap_user_exception, ) from marimo._runtime.executor import ( - ExecutionConfig, - get_executor, + Evaluator, + ExecutionLifecycle, + StrictLifecycle, + resolve_executor, ) from marimo._runtime.marimo_pdb import MarimoPdb from marimo._runtime.runner.hook_context import ( @@ -44,6 +39,8 @@ ExceptionOrError, ExecutionContextManager, ) +from marimo._runtime.runner.result import RunResult +from marimo._runtime.runner.scheduler import SequentialScheduler from marimo._sql.error_utils import ( create_sql_error_from_exception, is_sql_parse_error, @@ -53,7 +50,7 @@ LOGGER = marimo_logger() if TYPE_CHECKING: - from collections.abc import Iterator + from collections import deque from marimo._runtime.runner.hooks import NotebookCellHooks from marimo._runtime.state import State @@ -83,22 +80,7 @@ def cell_filename(cell_id: CellId_t) -> str: return f"" -@dataclass -class RunResult: - # Raw output of cell: last expression - output: Any - # Exception raised by cell, if any - # - # TODO(akshayka): Exceptions and "Errors" (most of which are at parse time - # and can't be encountered by the runner) shouldn't be packed into a single - # field. - exception: ExceptionOrError | None - # Accumulated output: via imperative mo.output.append() - accumulated_output: Any = None - - def success(self) -> bool: - """Whether the cell expected successfully""" - return self.exception is None +__all__ = ["RunResult", "Runner", "cell_filename", "should_show_traceback"] def should_show_traceback( @@ -137,9 +119,6 @@ def __init__( self.graph = graph self.debugger = debugger self.excluded_cells = excluded_cells or set() - self._executor = get_executor( - ExecutionConfig(is_strict=execution_type == "strict") - ) self.execution_context = execution_context self._hooks = hooks self.user_config = user_config @@ -155,27 +134,30 @@ def __init__( # so that they can be transitioned out of error if a future # run request repairs the graph self.roots = roots - self.cells_to_run: deque[CellId_t] = deque( - Runner.compute_cells_to_run( - self.graph, - self.roots, - self.excluded_cells, - self.execution_mode, - ) + cells_to_run_list = Runner.compute_cells_to_run( + self.graph, + self.roots, + self.excluded_cells, + self.execution_mode, ) - # tracks cancelled cells: raising cell -> descendants, with O(1) lookup - self.cancelled_cells = CancelledCells() - # whether the runner has been interrupted - self.interrupted = False + self._scheduler = SequentialScheduler(cells_to_run_list, self.graph) + # mapping from cell_id to exception it raised self.exceptions: dict[CellId_t, ExceptionOrError] = {} - # each cell's position in the run queue + # each cell's position in the original run queue self._run_position = { - cell_id: index for index, cell_id in enumerate(self.cells_to_run) + cell_id: index for index, cell_id in enumerate(cells_to_run_list) } + lifecycles: list[ExecutionLifecycle] = [] + if execution_type == "strict": + lifecycles.append(StrictLifecycle(self.graph)) + self._evaluator = Evaluator( + executor=resolve_executor(), lifecycles=lifecycles + ) + @staticmethod def compute_cells_to_run( graph: dataflow.DirectedGraph, @@ -219,71 +201,33 @@ def compute_cells_to_run( return sorted_cells - # Adapted from - # https://github.com/ipython/ipykernel/blob/eddd3e666a82ebec287168b0da7cfa03639a3772/ipykernel/ipkernel.py#L312 - @staticmethod - @contextlib.contextmanager - def _cancel_on_sigint(future: asyncio.Future[Any]) -> Iterator[None]: - """ContextManager for capturing SIGINT and cancelling a future + @property + def cells_to_run(self) -> deque[CellId_t]: + return self._scheduler.cells_to_run - SIGINT raises in the event loop when running async code, - but we want it to halt a coroutine. + @property + def cancelled_cells(self) -> CancelledCells: + return self._scheduler.cancelled_cells - Ideally, it would raise KeyboardInterrupt, but this turns it into a - CancelledError. - """ - sigint_future: asyncio.Future[int] = asyncio.Future() - - # whichever future finishes first, - # cancel the other one - def cancel_unless_done(f: asyncio.Future[Any], _: Any) -> None: - if f.cancelled() or f.done(): - return - f.cancel() - - # when sigint finishes, - # abort the coroutine with CancelledError - sigint_future.add_done_callback( - functools.partial(cancel_unless_done, future) - ) - # when the main future finishes, - # stop watching for SIGINT events - future.add_done_callback( - functools.partial(cancel_unless_done, sigint_future) - ) + @property + def interrupted(self) -> bool: + return self._scheduler.interrupted - def handle_sigint(*_: Any) -> None: - if sigint_future.cancelled() or sigint_future.done(): - return - # mark as done, to trigger cancellation - sigint_future.set_result(1) - - # set the custom sigint handler during this context - save_sigint = signal.signal(signal.SIGINT, handle_sigint) - try: - yield - finally: - # restore the previous sigint handler - signal.signal(signal.SIGINT, save_sigint) + @interrupted.setter + def interrupted(self, value: bool) -> None: + self._scheduler.interrupted = value def cancel(self, cell_id: CellId_t) -> None: """Mark a cell (and its descendants) as cancelled.""" - descendants = { - cid - for cid in dataflow.transitive_closure(self.graph, {cell_id}) - if cid in self.cells_to_run - } - self.cancelled_cells.add(cell_id, descendants) - for cid in descendants: - self.graph.cells[cid].set_run_result_status("cancelled") + self._scheduler.cancel(cell_id) def cancelled(self, cell_id: CellId_t) -> bool: """Return whether a cell has been cancelled.""" - return cell_id in self.cancelled_cells + return self._scheduler.cancelled(cell_id) def pending(self) -> bool: """Whether there are more cells to run.""" - return not self.interrupted and len(self.cells_to_run) > 0 + return self._scheduler.pending() def _get_run_position(self, cell_id: CellId_t) -> int | None: """Position in the original run queue""" @@ -357,7 +301,7 @@ def resolve_state_updates( def pop_cell(self) -> CellId_t: """Get the next cell to run.""" - return self.cells_to_run.popleft() + return self._scheduler.pop_cell() def _run_result_from_exception( self, @@ -461,65 +405,75 @@ async def run(self, cell_id: CellId_t) -> RunResult: self.debugger._last_traceback = None cell = self.graph.cells[cell_id] - run_result = None + # The Evaluator captures all body/lifecycle exceptions into the + # returned RunResult; cell_id-specific classification + side + # effects are applied below in `_finalize_run_result`. try: - if cell.is_coroutine(): - return_value_future = asyncio.ensure_future( - self._executor.execute_cell_async( - cell, - self.glbls, - self.graph, - ) - ) - if threading.current_thread() == threading.main_thread(): - # edit mode: need to handle user interrupts - with Runner._cancel_on_sigint(return_value_future): - return_value = await return_value_future - else: - # run mode: can't use signal.signal, not interruptible - # by user anyway. - return_value = await return_value_future - else: - return_value = self._executor.execute_cell( - cell, - self.glbls, - self.graph, - ) - run_result = RunResult(output=return_value, exception=None) - except asyncio.exceptions.CancelledError: - # User interrupt - # interrupt the entire runner - # Async cells can only be cancelled via a user interrupt - run_result = RunResult(output=None, exception=MarimoInterrupt()) - # Still provide a general traceback. - tmpio = io.StringIO() - traceback.print_exc(file=tmpio) - tmpio.seek(0) - write_traceback(tmpio.read()) - # Strict mode errors may also raise errors outside of execution. - except MarimoNameError as e: - self.cancel(cell_id) - strict_exception = MarimoStrictExecutionError(str(e), e.ref, None) - run_result = RunResult( - output=strict_exception, exception=strict_exception + raw_result = await self._evaluator.evaluate_interruptible( + cell, self.glbls ) - except MarimoMissingRefError as e: - # In strict mode, marimo refuses to evaluate a cell if there are - # missing definitions. Since the cell hasn't run, this is a pre - # check error, but still mark descendants as cancelled. + run_result = self._finalize_run_result(raw_result, cell_id) + except BaseException: + # Defensive: an unexpected escape from the Evaluator or a bug + # in `_finalize_run_result` would otherwise tear down the + # runner loop. Degrade gracefully with an empty RunResult. + LOGGER.error( + """marimo encountered an internal error. + + marimo finished executing a cell, but did not produce + a run result. + + Please copy this message and paste it in a GitHub issue: + + https://github.com/marimo-team/marimo/issues + + Any additional context of what caused this error, such + as sample code to reproduce, will help us debug. + """ + ) + run_result = RunResult(output=None, exception=None) + + # Mark as interrupted if the cell raised a MarimoInterrupt + # Set here since failed async can also trigger an Interrupt. + if isinstance(run_result.exception, MarimoInterrupt): + self.interrupted = True + + self._update_debugger_state(run_result, cell_id) + + if run_result.exception is not None: + self.exceptions[cell_id] = run_result.exception + + return run_result + + def _finalize_run_result( + self, raw_result: RunResult, cell_id: CellId_t + ) -> RunResult: + """Classify the Evaluator's RunResult and apply Runner side effects.""" + exc = raw_result.exception + if exc is None: + return raw_result + if not isinstance(exc, BaseException): + # No exception to handle. Cancel descendants and surface the payload + # as-is. self.cancel(cell_id) - ref, blamed_cell = self._get_blamed_cell(e) - name_output = MarimoStrictExecutionError( - "marimo was unable to resolve " - f"a reference to `{ref}` in cell : ", - ref, - blamed_cell, + return raw_result + + if isinstance(exc, asyncio.exceptions.CancelledError): + # Surface cancellation as a MarimoInterrupt for downstream handling. + tmpio = io.StringIO() + traceback.print_exception( + type(exc), exc, exc.__traceback__, file=tmpio ) - run_result = RunResult(output=name_output, exception=name_output) + tmpio.seek(0) + write_traceback(tmpio.read()) + return RunResult(output=None, exception=MarimoInterrupt()) + # Should cover all cell runtime exceptions. - except MarimoRuntimeException as e: - output: Any = None - unwrapped_exception: BaseException | None = e.__cause__ + if isinstance(exc, MarimoRuntimeException): + # Unwrap the user exception and upgrade a raw NameError to + # MarimoMissingRefError when the missing name is defined + # elsewhere in the graph. + unwrapped_exception = unwrap_user_exception(exc, self.graph) # Interrupts are sometimes sent multiple times; in particular, # it appears that polars forwards interrupts, so interrupting @@ -531,7 +485,7 @@ async def run(self, cell_id: CellId_t) -> RunResult: try: run_result, unwrapped_exception = ( self._run_result_from_exception( - output, unwrapped_exception, cell_id + None, unwrapped_exception, cell_id ) ) except KeyboardInterrupt: @@ -541,8 +495,8 @@ async def run(self, cell_id: CellId_t) -> RunResult: # Exceptions trigger cancellation of descendants. # - # TODO(akshayka): Another interrupt will end up interrupting - # this call as well, so this should be lifted out of `run`. + # TODO(akshayka): A SIGINT during cancel() can interrupt this + # call, so this should be lifted to a non-interruptible path. self.cancel(cell_id) if should_show_traceback(run_result.exception): @@ -567,81 +521,54 @@ async def run(self, cell_id: CellId_t) -> RunResult: ) tmpio.seek(0) write_traceback(tmpio.read()) - except BaseException as e: - # Check that MarimoRuntimeException has't already handled the - # error, since exceptions fall through except blocks. - # If not, then this is an unexpected error. - if not isinstance(e, MarimoRuntimeException): - LOGGER.error(f"Unexpected error type: {e}") - self.cancel(cell_id) - unknown_error = UnknownError(f"{e}") - run_result = RunResult(output=None, exception=unknown_error) - tmpio = io.StringIO() - traceback.print_exc(file=tmpio) - tmpio.seek(0) - write_traceback(tmpio.read()) - finally: - # TODO(akshayka): some of this logic should be lifted out - # of `run`, (in particular to where execution context is not set) - # so that it is not interruptible - if run_result is None: - LOGGER.error( - """marimo encountered an internal error. - - marimo finished executing a cell, but did not produce - a run result. - - Please copy this message and paste it in a GitHub issue: - - https://github.com/marimo-team/marimo/issues - - Any additional context of what caused this error, such - as sample code to reproduce, will help us debug. - """ - ) - run_result = RunResult(output=None, exception=None) - - # Mark as interrupted if the cell raised a MarimoInterrupt - # Set here since failed async can also trigger an Interrupt. - if isinstance(run_result.exception, MarimoInterrupt): - self.interrupted = True - - # if a debugger is active, force it to skip past marimo code. - try: - # Bdb defines the botframe attribute and sets it to non-None - # when it starts up - if self.debugger is not None: - if ( - hasattr(self.debugger, "botframe") - and self.debugger.botframe is not None - ): - self.debugger.set_continue() - # Hold on to this information for debugging postmortem etc. - if run_result.exception is not None and hasattr( - run_result.exception, "__traceback__" - ): - tb = run_result.exception.__traceback__ - if isinstance(tb, TracebackType): - self.debugger._last_traceback = tb - self.debugger._last_tracebacks[cell_id] = tb - except Exception as debugger_error: - # This has never been hit, but just in case -- don't want - # to crash the kernel. - LOGGER.error( - """Internal marimo error. Please copy this message and - paste it in a GitHub issue: - - https://github.com/marimo-team/marimo/issues - - An exception raised attempting to continue debugger (%s). - """, - str(debugger_error), - ) - - if run_result.exception is not None: - self.exceptions[cell_id] = run_result.exception - - return run_result + return run_result + + # Anything else escaping the Evaluator is unexpected. + LOGGER.error(f"Unexpected error type: {exc}") + self.cancel(cell_id) + tmpio = io.StringIO() + traceback.print_exception( + type(exc), exc, exc.__traceback__, file=tmpio + ) + tmpio.seek(0) + write_traceback(tmpio.read()) + return RunResult(output=None, exception=UnknownError(f"{exc}")) + + def _update_debugger_state( + self, run_result: RunResult, cell_id: CellId_t + ) -> None: + """Skip marimo frames in the debugger and stash the cell's traceback.""" + # if a debugger is active, force it to skip past marimo code. + try: + # Bdb defines the botframe attribute and sets it to non-None + # when it starts up + if self.debugger is not None: + if ( + hasattr(self.debugger, "botframe") + and self.debugger.botframe is not None + ): + self.debugger.set_continue() + # Hold on to this information for debugging postmortem etc. + if run_result.exception is not None and hasattr( + run_result.exception, "__traceback__" + ): + tb = run_result.exception.__traceback__ + if isinstance(tb, TracebackType): + self.debugger._last_traceback = tb + self.debugger._last_tracebacks[cell_id] = tb + except Exception as debugger_error: + # This has never been hit, but just in case -- don't want + # to crash the kernel. + LOGGER.error( + """Internal marimo error. Please copy this message and + paste it in a GitHub issue: + + https://github.com/marimo-team/marimo/issues + + An exception raised attempting to continue debugger (%s). + """, + str(debugger_error), + ) def _get_blamed_cell( self, e: MarimoMissingRefError diff --git a/marimo/_runtime/runner/hooks_post_execution.py b/marimo/_runtime/runner/hooks_post_execution.py index 8b82bf0ede5..6efefc29de0 100644 --- a/marimo/_runtime/runner/hooks_post_execution.py +++ b/marimo/_runtime/runner/hooks_post_execution.py @@ -35,7 +35,6 @@ ) from marimo._messaging.tracebacks import ( _highlight_traceback, - _trim_traceback, format_exception_message, write_traceback, ) @@ -430,9 +429,7 @@ def _broadcast_outputs( and run_result.exception.__traceback__ ): tb_lines = tb.format_exception(run_result.exception) - formatted_traceback = _highlight_traceback( - _trim_traceback("".join(tb_lines)) - ) + formatted_traceback = _highlight_traceback("".join(tb_lines)) CellNotificationUtils.broadcast_error( data=[ diff --git a/marimo/_runtime/runner/result.py b/marimo/_runtime/runner/result.py new file mode 100644 index 00000000000..0be0f17a3b3 --- /dev/null +++ b/marimo/_runtime/runner/result.py @@ -0,0 +1,28 @@ +# Copyright 2026 Marimo. All rights reserved. +"""The value type a cell produces when it runs.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from marimo._runtime.runner.hook_context import ExceptionOrError + + +@dataclass +class RunResult: + # Raw output of cell: last expression + output: Any + # Exception raised by cell, if any + # + # TODO(akshayka): Exceptions and "Errors" (most of which are at parse time + # and can't be encountered by the runner) shouldn't be packed into a single + # field. + exception: ExceptionOrError | None + # Accumulated output: via imperative mo.output.append() + accumulated_output: Any = None + + def success(self) -> bool: + """Whether the cell executed successfully""" + return self.exception is None diff --git a/marimo/_runtime/runner/scheduler.py b/marimo/_runtime/runner/scheduler.py new file mode 100644 index 00000000000..6205d2670dc --- /dev/null +++ b/marimo/_runtime/runner/scheduler.py @@ -0,0 +1,89 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Scheduler owns the cell queue and cancellation state""" + +from __future__ import annotations + +from collections import deque +from typing import TYPE_CHECKING, Protocol + +from marimo._runtime import dataflow +from marimo._runtime.runner.hook_context import CancelledCells + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator, Sequence + + from marimo._runtime.dataflow import DirectedGraph + from marimo._types.ids import CellId_t + + +class Scheduler(Protocol): + """Cell queue + cancellation. Surface for future scheduler types.""" + + def pending(self) -> bool: ... + def pop_cell(self) -> CellId_t: ... + def cancel(self, cell_id: CellId_t) -> None: ... + def cancelled(self, cell_id: CellId_t) -> bool: ... + def batch( + self, cell_ids: Iterable[CellId_t] + ) -> Iterator[list[CellId_t]]: ... + + +class SequentialScheduler: + """Single-threaded FIFO queue + cancellation.""" + + def __init__( + self, + cells_to_run: Sequence[CellId_t], + graph: DirectedGraph, + ) -> None: + self._cells_to_run: deque[CellId_t] = deque(cells_to_run) + self._cancelled = CancelledCells() + self._graph = graph + self._interrupted = False + + def pending(self) -> bool: + return not self._interrupted and len(self._cells_to_run) > 0 + + def pop_cell(self) -> CellId_t: + return self._cells_to_run.popleft() + + def batch(self, cell_ids: Iterable[CellId_t]) -> Iterator[list[CellId_t]]: + """Yield batches of cells to execute. + + Sequential default: one cell per batch. + """ + self._cells_to_run.clear() + self._cells_to_run.extend(cell_ids) + while self._cells_to_run and not self._interrupted: + yield [self._cells_to_run.popleft()] + + def cancel(self, cell_id: CellId_t) -> None: + """Mark a cell and its descendants as cancelled.""" + descendants = { + cid + for cid in dataflow.transitive_closure(self._graph, {cell_id}) + if cid in self._cells_to_run + } + self._cancelled.add(cell_id, descendants) + for cid in descendants: + self._graph.cells[cid].set_run_result_status("cancelled") + + def cancelled(self, cell_id: CellId_t) -> bool: + return cell_id in self._cancelled + + @property + def interrupted(self) -> bool: + return self._interrupted + + @interrupted.setter + def interrupted(self, value: bool) -> None: + self._interrupted = value + + @property + def cancelled_cells(self) -> CancelledCells: + return self._cancelled + + @property + def cells_to_run(self) -> deque[CellId_t]: + """The live queue. Mutates as cells are popped.""" + return self._cells_to_run diff --git a/marimo/_runtime/runtime.py b/marimo/_runtime/runtime.py index 0eacd1d26c7..7b21537a45e 100644 --- a/marimo/_runtime/runtime.py +++ b/marimo/_runtime/runtime.py @@ -479,8 +479,6 @@ def __init__( # timestamp, to save the user from having to spam the interrupt button self.last_interrupt_timestamp: float | None = None - # Named attributes exist because internal kernel paths (run hooks, - # script metadata) and tests reach into specific callbacks directly. self.secrets_callbacks = SecretsCallbacks(self) self.datasets_callbacks = DatasetCallbacks(self) self.packages_callbacks = PackagesCallbacks(self) diff --git a/marimo/_types/globals.py b/marimo/_types/globals.py new file mode 100644 index 00000000000..ab51826c4f7 --- /dev/null +++ b/marimo/_types/globals.py @@ -0,0 +1,15 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Type aliases for cell globals dicts. + +`MutableGlobals` is the concrete `dict` passed through `exec` / +`eval`; `Globals` is the read-only view for consumers that only +inspect the dict (e.g. collecting a cell's defs after execution). +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, TypeAlias + +Globals: TypeAlias = Mapping[str, Any] +MutableGlobals: TypeAlias = dict[str, Any] diff --git a/tests/_ast/test_app.py b/tests/_ast/test_app.py index 69105a29146..ec9ec8482db 100644 --- a/tests/_ast/test_app.py +++ b/tests/_ast/test_app.py @@ -1,15 +1,13 @@ # Copyright 2026 Marimo. All rights reserved. from __future__ import annotations -import os + import pathlib import subprocess import sys import textwrap from typing import TYPE_CHECKING, Any -from unittest.mock import patch -import click import pytest from marimo._ast.app import ( @@ -28,21 +26,13 @@ SetupRootError, UnparsableError, ) -from marimo._ast.load import load_app from marimo._ast.names import SETUP_CELL_NAME -from marimo._convert.converters import MarimoConvert from marimo._dependencies.dependencies import DependencyManager from marimo._plugins.stateless.flex import vstack +from marimo._runtime.commands import UpdateUIElementCommand from marimo._runtime.context.types import get_context -from marimo._runtime.commands import UpdateUIElementCommand, ExecuteCellCommand -from marimo._schemas.serialization import ( - AppInstantiation, - CellDef, - NotebookSerializationV1, -) from marimo._types.ids import CellId_t -from tests.conftest import ExecReqProvider, MockedKernel -from tests._messaging.mocks import MockStream +from tests.conftest import ExecReqProvider if TYPE_CHECKING: from marimo._runtime.runtime import Kernel @@ -102,7 +92,7 @@ def test_run_with_docstring() -> None: @app.cell def _() -> tuple[object]: - doc = __doc__ # noqa: F821 + doc = __doc__ return (doc,) _, defs = app.run() @@ -115,7 +105,7 @@ def test_run_with_no_docstring() -> None: @app.cell def _() -> tuple[object]: - doc = __doc__ # noqa: F821 + doc = __doc__ return (doc,) _, defs = app.run() @@ -133,7 +123,9 @@ def config() -> tuple[int, float]: return batch_size, learning_rate @app.cell - def process_data(batch_size: int, learning_rate: float) -> tuple[float]: + def process_data( + batch_size: int, learning_rate: float + ) -> tuple[float]: result = batch_size * learning_rate return (result,) @@ -150,7 +142,9 @@ def other_cell() -> tuple[str]: assert defs["message"] == "independent" # Test 2: Run with overridden values - outputs, defs = app.run(defs={"batch_size": 64, "learning_rate": 0.001}) + outputs, defs = app.run( + defs={"batch_size": 64, "learning_rate": 0.001} + ) assert defs["batch_size"] == 64 assert defs["learning_rate"] == 0.001 assert defs["result"] == 64 * 0.001 @@ -203,7 +197,6 @@ def test_run_with_refs_setup_cell_protection() -> None: app = App() with app.setup: - import os setup_var = "from_setup" @app.cell @@ -232,7 +225,6 @@ def test_run_with_undefined_refs_in_setup_cell() -> None: app = App() with app.setup: - import os a = 1 if a > 2: setup_var = "from_setup" @@ -248,7 +240,6 @@ def use_setup(setup_var: str) -> tuple[str]: app.run() assert "setup_var" in str(exc_info.value) - @staticmethod def test_setup() -> None: app = App() @@ -291,7 +282,6 @@ def __() -> tuple[int, int]: assert (defs["y"], defs["z"]) == (1, 2) assert defs["a"] == 2 - @staticmethod def test_cycle() -> None: app = App() @@ -315,11 +305,11 @@ def test_cycle_missing_args_rets() -> None: @app.cell def one() -> None: - x = y # noqa: F841, F821 + x = y # noqa: F821 @app.cell def two() -> None: - y = x # noqa: F841, F821 + y = x # noqa: F821 with pytest.raises(CycleError): app.run() @@ -347,11 +337,11 @@ def test_multiple_definitions_missing_args_rets() -> None: @app.cell def one() -> None: - x = 0 # noqa: F841 + x = 0 @app.cell def two() -> None: - x = 0 # noqa: F841 + x = 0 with pytest.raises(MultipleDefinitionError): app.run() @@ -362,11 +352,11 @@ def test_delete_nonlocal_ok() -> None: @app.cell def one() -> None: - x = 0 # noqa: F841 + x = 0 @app.cell def two() -> None: - del x # noqa: F841, F821 + del x # noqa: F821 # smoke test, no error raised app.run() @@ -426,7 +416,7 @@ def test_resolve_var_not_local_from_nested_scope() -> None: @app.cell def _() -> tuple[str]: - _x = 10 # noqa: F841 + _x = 10 def _f() -> str: _x = "nested" @@ -454,7 +444,7 @@ def _f() -> str: @app.cell def _() -> None: - _x # type: ignore # noqa: F821 + _x # type: ignore return with pytest.raises(NameError) as e: @@ -468,7 +458,7 @@ def test_locals_dont_leak() -> None: @app.cell def _() -> None: - _x = 0 # noqa: F841 + _x = 0 return @app.cell @@ -503,7 +493,7 @@ def test_dunder_rewritten_as_local() -> None: @app.cell def _() -> None: - __ = 1 # noqa: F841 + __ = 1 return @app.cell @@ -648,6 +638,7 @@ def test_run_mo_stop() -> None: @app.cell def _() -> Any: import marimo as mo + return (mo,) @app.cell @@ -672,6 +663,7 @@ def test_run_mo_stop_descendant() -> None: @app.cell def _() -> Any: import marimo as mo + return (mo,) @app.cell @@ -697,6 +689,7 @@ def test_run_mo_stop_descendant_multiple() -> None: @app.cell def _() -> Any: import marimo as mo + return (mo,) @app.cell @@ -711,7 +704,6 @@ def _(mo) -> tuple[int]: y = 0 return (y,) - @app.cell def _(x) -> tuple[int]: x @@ -724,13 +716,42 @@ def _(y) -> tuple[int]: b = 0 return - _, defs = app.run() assert "x" not in defs assert "y" not in defs assert "a" not in defs assert "b" not in defs + @staticmethod + def test_run_mo_stop_records_output() -> None: + # mo.stop's `output=` arg is shown to the user in edit/kernel mode. + # Script mode used to silently drop it; cell_runner has always + # recorded it. The runner-consolidation refactor aligns these. + app = App() + + @app.cell + def first() -> Any: + import marimo as mo + + return (mo,) + + @app.cell + def stop_cell(mo) -> tuple[int]: + mo.stop(True, output="stopped-output-value") + x = 0 + return (x,) + + @app.cell + def descendant(x) -> tuple[int]: + y = x + 1 + return (y,) + + outputs, defs = app.run() + # stop_cell at index 1 records the stop's output; its descendant + # is cancelled and absent from the flattened output tuple. + assert outputs == (None, "stopped-output-value") + assert "x" not in defs + assert "y" not in defs @staticmethod def test_run_mo_stop_async() -> None: @@ -739,6 +760,7 @@ def test_run_mo_stop_async() -> None: @app.cell def _() -> Any: import marimo as mo + return (mo,) @app.cell @@ -763,6 +785,7 @@ def test_run_mo_stop_descendant_async() -> None: @app.cell def _() -> Any: import marimo as mo + return (mo,) @app.cell @@ -781,7 +804,6 @@ async def _(x) -> tuple[int]: assert "x" not in defs assert "y" not in defs - @pytest.mark.skipif( condition=not DependencyManager.matplotlib.has(), reason="requires matplotlib", @@ -955,11 +977,16 @@ def __(): # Public mutable fields should be deep-copied, not shared assert original_impl.config is not cloned_impl.config - assert original_impl.import_workspace is not cloned_impl.import_workspace + assert ( + original_impl.import_workspace is not cloned_impl.import_workspace + ) # Private mutable runtime state fields should also be independent assert original_impl._status is not cloned_impl._status - assert original_impl._run_result_status is not cloned_impl._run_result_status + assert ( + original_impl._run_result_status + is not cloned_impl._run_result_status + ) assert original_impl._stale is not cloned_impl._stale assert original_impl._output is not cloned_impl._output @@ -1030,8 +1057,7 @@ class TestInvalidSetup: @staticmethod def test_initial_setup() -> None: app = App() - app._unparsable_cell(";", - name="setup") + app._unparsable_cell(";", name="setup") assert app._cell_manager.has_cell("setup") assert app._cell_manager.cell_name("setup") == "setup" @@ -1039,22 +1065,21 @@ def test_initial_setup() -> None: @staticmethod def test_not_initial_setup() -> None: app = App() - app._unparsable_cell(";", - name="other") - app._unparsable_cell(";", - name="setup") + app._unparsable_cell(";", name="other") + app._unparsable_cell(";", name="setup") assert not app._cell_manager.has_cell("setup") @staticmethod def test_not_initial_setup_cell() -> None: app = App() + @app.cell def _(): def B() -> float: return 1.0 - app._unparsable_cell(";", - name="setup") + + app._unparsable_cell(";", name="setup") assert not app._cell_manager.has_cell("setup") @@ -1205,7 +1230,9 @@ def __() -> tuple[int]: with pytest.raises(ValueError) as excinfo: await app.embed(defs={"x": mo.ui.slider(1, 10)}) - assert "Substituting UI Elements for variables is not allowed" in str(excinfo.value) + assert "Substituting UI Elements for variables is not allowed" in str( + excinfo.value + ) async def test_app_embed_with_defs_multiple_vars(self) -> None: """Test embed() with defs overriding a cell that defines multiple variables.""" @@ -1532,8 +1559,6 @@ def test_app_not_changed() -> None: with app.setup: app = 1 - - @staticmethod def test_setup_not_exposed() -> None: app = App() @@ -1545,7 +1570,6 @@ def test_setup_not_exposed() -> None: except NameError: x = False - @staticmethod def test_setup_in_memory() -> None: app = App() @@ -1598,11 +1622,8 @@ def test_setup_hide_code() -> None: assert setup_cell is not None assert setup_cell.config.hide_code is False - @staticmethod - async def test_app_embed_preserves_file_path( - app: App - ) -> None: + async def test_app_embed_preserves_file_path(app: App) -> None: with app.setup: from tests._ast.app_data import notebook_filename @@ -1624,7 +1645,6 @@ def _(cloned: AppEmbedResult, filename: str, directory: str) -> None: assert cloned.defs.get("this_is_foo_file").endswith(filename) assert cloned.defs.get("this_is_foo_path").stem == directory - @staticmethod async def test_app_embed_in_kernel( k: Kernel, exec_req: ExecReqProvider @@ -1648,10 +1668,13 @@ async def test_app_embed_in_kernel( filename = "notebook_filename.py" directory = "app_data" assert k.globals["app"].defs.get("this_is_foo_file").endswith(filename) - assert k.globals["cloned"].defs.get("this_is_foo_file").endswith(filename) + assert ( + k.globals["cloned"].defs.get("this_is_foo_file").endswith(filename) + ) assert k.globals["app"].defs.get("this_is_foo_path").stem == directory - assert k.globals["cloned"].defs.get("this_is_foo_path").stem == directory - + assert ( + k.globals["cloned"].defs.get("this_is_foo_path").stem == directory + ) @staticmethod async def test_app_embed_same_cell_in_kernel( @@ -1693,8 +1716,9 @@ async def test_imported_app_has_prefixed_setup_cell( This tests the fix where setup cells get the prefix like other cells. """ - await k.run([ - exec_req.get(""" + await k.run( + [ + exec_req.get(""" # Import in kernel context; the prefix the app gets # depends on whether it was first imported in a kernel context, # so we reload it in case notebook_filename was loaded elsewhere @@ -1705,11 +1729,14 @@ async def test_imported_app_has_prefixed_setup_cell( importlib.reload(mod) app = mod.app """) - ]) + ] + ) assert not k.errors nb_app = k.globals["app"] cell_ids = list(InternalApp(nb_app).cell_manager.cell_ids()) - setup_cell_ids = [cid for cid in cell_ids if cid.endswith(SETUP_CELL_NAME)] + setup_cell_ids = [ + cid for cid in cell_ids if cid.endswith(SETUP_CELL_NAME) + ] assert len(setup_cell_ids) == 1 assert is_external_cell_id(setup_cell_ids[0]) @@ -1780,9 +1807,7 @@ def _(): internal_app = InternalApp(app) cell_id = next(iter(internal_app.cell_manager.cell_ids())) - original_compiled = internal_app.cell_manager._compiled_cells[ - cell_id - ] + original_compiled = internal_app.cell_manager._compiled_cells[cell_id] assert original_compiled is not None internal_app.with_data( @@ -1904,7 +1929,6 @@ def __(x: int) -> tuple[int]: assert not k.errors assert k.globals["overrides"] == {"x": 100} - @pytest.mark.xfail( True, reason="Flaky in CI, can't repro locally", strict=False ) diff --git a/tests/_messaging/test_tracebacks.py b/tests/_messaging/test_tracebacks.py index 1346eba1957..f5c5de5d4c7 100644 --- a/tests/_messaging/test_tracebacks.py +++ b/tests/_messaging/test_tracebacks.py @@ -9,7 +9,6 @@ from marimo._messaging.context import HTTP_REQUEST_CTX, is_code_mode_request from marimo._messaging.tracebacks import ( _highlight_traceback, - _trim_traceback, format_exception_message, is_code_highlighting, write_traceback, @@ -259,11 +258,3 @@ def test_empty_url(self) -> None: assert is_code_mode_request() is False finally: HTTP_REQUEST_CTX.reset(token) - - def test_trim(self) -> None: - prefix = "Traceback (most recent call last):\n" - head = ' File ".../marimo/_runtime/executor.py", line 139, in execute_cell\n return eval(cell.last_expr, glbls)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n' - rest = ( - ' File ".../__marimo__cell_Hbol_.py", line 2, in \n...\n' - ) - assert _trim_traceback(f"{prefix}{head}{rest}") == f"{prefix}{rest}" diff --git a/tests/_runtime/runner/test_cell_runner.py b/tests/_runtime/runner/test_cell_runner.py index 9da65671c89..4ab2f601517 100644 --- a/tests/_runtime/runner/test_cell_runner.py +++ b/tests/_runtime/runner/test_cell_runner.py @@ -276,3 +276,132 @@ async def test_converging_runs_when_all_branches_trigger( assert "b" in k.globals assert "result" in k.globals assert k.graph.cells["res"].run_result_status == "success" + + +# --- Surface 3: registered plugin Executor runs via Runner ------------------ + + +async def test_runner_dispatches_to_registered_plugin_executor( + execution_kernel: Kernel, + exec_req: ExecReqProvider, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A factory registered against `marimo.cell.executor` is the one + the kernel `Runner` dispatches through.""" + from typing import Any + + from marimo._runtime.executor.evaluator import _EXECUTOR_REGISTRY + + recorded: list[str] = [] + sentinel_output = object() + + class _SentinelExecutor: + name = "sentinel" + + def execute_cell(self, cell: Any, glbls: dict[str, Any]) -> object: + del glbls + recorded.append(cell.cell_id) + return sentinel_output + + async def execute_cell_async( + self, cell: Any, glbls: dict[str, Any] + ) -> object: + del glbls + recorded.append(cell.cell_id) + return sentinel_output + + def factory() -> _SentinelExecutor: + return _SentinelExecutor() + + # Populate the kernel first (uses the real DefaultExecutor — the + # registry isn't patched yet). + k = execution_kernel + await k.run([er := exec_req.get("'hello'; 123")]) + + # Fully isolate the registry: replace both `_plugins` and + # `names` so installed third-party entry points can't shadow the + # sentinel. monkeypatch restores both on teardown. + monkeypatch.setattr(_EXECUTOR_REGISTRY, "_plugins", {"sentinel": factory}) + monkeypatch.setattr(_EXECUTOR_REGISTRY, "names", lambda: ["sentinel"]) + + runner = Runner( + roots=set(k.graph.cells.keys()), + graph=k.graph, + glbls=k.globals, + debugger=k.debugger, + hooks=NotebookCellHooks(), + ) + run_result = await runner.run(er.cell_id) + + assert recorded == [er.cell_id] + assert run_result.output is sentinel_output + + +# --- Surface 4: Runner.interrupted flips on cancellation -------------------- + + +async def test_runner_interrupted_flag_flips_on_sync_marimo_interrupt( + execution_kernel: Kernel, exec_req: ExecReqProvider +) -> None: + """Sync cell body raising `MarimoInterrupt` (== `KeyboardInterrupt`) + surfaces as a bare `MarimoInterrupt` in the run result and flips + `runner.interrupted`.""" + k = execution_kernel + await k.run([er := exec_req.get("raise KeyboardInterrupt")]) + + runner = Runner( + roots=set(k.graph.cells.keys()), + graph=k.graph, + glbls=k.globals, + debugger=k.debugger, + hooks=NotebookCellHooks(), + ) + with capture_stderr(): + await runner.run(er.cell_id) + + assert runner.interrupted is True + + +async def test_runner_interrupted_flag_flips_on_async_cell_cancellation( + execution_kernel: Kernel, + exec_req: ExecReqProvider, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An async cell cancelled mid-await flips `runner.interrupted`. + + A bare `asyncio.CancelledError` arriving in `RunResult.exception` is + converted to `MarimoInterrupt` by the bare-`CancelledError` branch + of `_finalize_run_result`, which `run()` recognises to flip the + flag. + + Simulates the evaluator output directly (a bare `CancelledError` in + the `RunResult`) so this test is independent of the executor's + coroutine compilation. + """ + import asyncio + + from marimo._runtime.runner.result import RunResult + + k = execution_kernel + await k.run([er := exec_req.get("123")]) + + runner = Runner( + roots=set(k.graph.cells.keys()), + graph=k.graph, + glbls=k.globals, + debugger=k.debugger, + hooks=NotebookCellHooks(), + ) + + async def fake_evaluate(cell, glbls): # type: ignore[no-untyped-def] + del cell, glbls + return RunResult(output=None, exception=asyncio.CancelledError()) + + monkeypatch.setattr( + runner._evaluator, "evaluate_interruptible", fake_evaluate + ) + + with capture_stderr(): + await runner.run(er.cell_id) + + assert runner.interrupted is True diff --git a/tests/_runtime/test_dataflow.py b/tests/_runtime/test_dataflow.py index e220e576821..0f54e792161 100644 --- a/tests/_runtime/test_dataflow.py +++ b/tests/_runtime/test_dataflow.py @@ -10,6 +10,8 @@ from marimo._ast.visitor import Name, VariableData from marimo._dependencies.dependencies import DependencyManager from marimo._runtime import dataflow +from marimo._runtime.runner import by_refs +from marimo._runtime.runner.by_refs import _get_ancestors parse_cell = partial(compiler.compile_cell, cell_id="0") @@ -944,7 +946,7 @@ def test_is_disabled() -> None: def test_runner_sync() -> None: - """Test the Runner class for synchronous execution.""" + """Synchronous Cell.run(**kwargs) path. Must work without an event loop.""" graph = dataflow.DirectedGraph() # Create a chain of cells: 0 -> 1 -> 2 @@ -960,18 +962,15 @@ def test_runner_sync() -> None: third_cell = compiler.compile_cell(code, cell_id="2") graph.register_cell("2", third_cell) - # Create a runner - runner = dataflow.Runner(graph) - # Run the last cell - output, defs = runner.run_cell_sync("2", {}) + output, defs = by_refs.run_cell_sync(graph, "2", {}) # Check output and definitions assert output == 25 # 10 * 2 + 5 assert defs == {"z": 25} # Run the last cell with substituted values - output, defs = runner.run_cell_sync("2", {"y": 50}) + output, defs = by_refs.run_cell_sync(graph, "2", {"y": 50}) # Check output and definitions with substituted value assert output == 55 # 50 + 5 @@ -979,14 +978,14 @@ def test_runner_sync() -> None: # Try to run with an invalid argument try: - runner.run_cell_sync("2", {"invalid": 100}) + by_refs.run_cell_sync(graph, "2", {"invalid": 100}) raise AssertionError("Should have raised an exception") except ValueError: pass # Expected def test_runner_ancestors() -> None: - """Test that the Runner correctly identifies ancestors based on refs.""" + """Ancestor pruning based on substituted refs.""" graph = dataflow.DirectedGraph() # Create cells with different refs/defs patterns @@ -1002,19 +1001,16 @@ def test_runner_ancestors() -> None: third_cell = compiler.compile_cell(code, cell_id="2") graph.register_cell("2", third_cell) - # Create a runner - runner = dataflow.Runner(graph) - # Get ancestors of the third cell - ancestors = runner._get_ancestors(graph.cells["2"], {}) + ancestors = _get_ancestors(graph, graph.cells["2"], {}) assert ancestors == {"0", "1"} # When substituting y, only cell 0 should be an ancestor - ancestors = runner._get_ancestors(graph.cells["2"], {"y": 30}) + ancestors = _get_ancestors(graph, graph.cells["2"], {"y": 30}) assert ancestors == {"0"} # When substituting both x and y, there should be no ancestors - ancestors = runner._get_ancestors(graph.cells["2"], {"x": 40, "y": 30}) + ancestors = _get_ancestors(graph, graph.cells["2"], {"x": 40, "y": 30}) assert ancestors == set() diff --git a/tests/_runtime/test_executor_evaluator.py b/tests/_runtime/test_executor_evaluator.py new file mode 100644 index 00000000000..6e05053b917 --- /dev/null +++ b/tests/_runtime/test_executor_evaluator.py @@ -0,0 +1,687 @@ +# Copyright 2026 Marimo. All rights reserved. +# Stub classes here conform to the ExecutionLifecycle / Executor +# Protocols, so they take `cell` / `glbls` even when the test body +# doesn't use them. +# ruff: noqa: ARG001, ARG002 +"""Tests for the Evaluator + ExecutionLifecycle composition. + +Covers setup chain order, Skip termination, teardown reverse order, +teardown visibility of body exceptions, teardown-wins semantics on +double raise, and KeyboardInterrupt propagation through teardown. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from marimo._runtime.exceptions import MarimoRuntimeException +from marimo._runtime.executor import ( + DefaultExecutor, + Evaluator, + ExecutionLifecycle, + Skip, +) +from marimo._runtime.runner.result import RunResult + + +class _Recorder: + """Lifecycle that records setup/teardown calls into a shared log.""" + + def __init__( + self, + log: list[str], + tag: str, + skip: Skip | None = None, + setup_raises: BaseException | None = None, + teardown_raises: BaseException | None = None, + ) -> None: + self.name = f"recorder-{tag}" + self._log = log + self._tag = tag + self._skip = skip + self._setup_raises = setup_raises + self._teardown_raises = teardown_raises + self.last_run_result: Any = None + + def setup(self, cell: Any, glbls: dict[str, Any]) -> Skip | None: + self._log.append(f"setup:{self._tag}") + if self._setup_raises is not None: + raise self._setup_raises + return self._skip + + def teardown( + self, cell: Any, glbls: dict[str, Any], run_result: Any + ) -> None: + self._log.append(f"teardown:{self._tag}") + self.last_run_result = run_result + if self._teardown_raises is not None: + raise self._teardown_raises + + +class _StubExecutor: + """Executor that runs a caller-provided body, no exec/eval.""" + + name = "stub" + + def __init__(self, body: Any) -> None: + self._body = body + + def execute_cell(self, cell: Any, glbls: dict[str, Any]) -> Any: + return self._body(cell, glbls) + + async def execute_cell_async( + self, cell: Any, glbls: dict[str, Any] + ) -> Any: + result = self._body(cell, glbls) + if asyncio.iscoroutine(result): + return await result + return result + + +async def test_skip_terminates_setup_chain_but_runs_completed_teardowns() -> ( + None +): + log: list[str] = [] + a = _Recorder( + log, "A", skip=Skip(result=RunResult(output=42, exception=None)) + ) + b = _Recorder(log, "B") + + body_ran = [False] + + def body(cell: Any, glbls: dict[str, Any]) -> Any: + body_ran[0] = True + return "should-not-see-this" + + ev = Evaluator(executor=_StubExecutor(body), lifecycles=[a, b]) + result = await ev.evaluate(cell=None, glbls={}) + + assert result.output == 42 + assert result.exception is None + assert body_ran[0] is False + # A setup ran, A teardown ran. B setup did NOT run, B teardown did + # NOT run. + assert log == ["setup:A", "teardown:A"] + + +async def test_skip_result_preserves_accumulated_output() -> None: + """`Skip(result=RunResult(...))` threads the entire RunResult + through teardown — `output`, `exception`, and + `accumulated_output` all survive, including any future fields + added to `RunResult`.""" + log: list[str] = [] + skip_result = RunResult( + output="cached", exception=None, accumulated_output="streamed" + ) + a = _Recorder(log, "A", skip=Skip(result=skip_result)) + + ev = Evaluator(executor=_StubExecutor(lambda *_: "unused"), lifecycles=[a]) + result = await ev.evaluate(cell=None, glbls={}) + + assert result.output == "cached" + assert result.accumulated_output == "streamed" + assert result.exception is None + # Teardown saw the same RunResult that came back out. + assert a.last_run_result is result + + +async def test_teardowns_fire_in_reverse_order_on_success() -> None: + log: list[str] = [] + a = _Recorder(log, "A") + b = _Recorder(log, "B") + c = _Recorder(log, "C") + + ev = Evaluator( + executor=_StubExecutor(lambda *_: "ok"), + lifecycles=[a, b, c], + ) + result = await ev.evaluate(cell=None, glbls={}) + + assert result.output == "ok" + assert result.exception is None + assert log == [ + "setup:A", + "setup:B", + "setup:C", + "teardown:C", + "teardown:B", + "teardown:A", + ] + + +async def test_teardown_sees_body_exception_via_run_result() -> None: + log: list[str] = [] + a = _Recorder(log, "A") + + def boom(cell: Any, glbls: dict[str, Any]) -> Any: + raise ValueError("body bomb") + + ev = Evaluator(executor=_StubExecutor(boom), lifecycles=[a]) + # The _StubExecutor doesn't wrap user exceptions; the body's + # ValueError lands directly in result.exception, and the teardown + # sees that same exception via run_result. + result = await ev.evaluate(cell=None, glbls={}) + + assert isinstance(result.exception, ValueError) + assert str(result.exception) == "body bomb" + assert a.last_run_result is not None + assert isinstance(a.last_run_result.exception, ValueError) + + +async def test_default_executor_wraps_user_exception_in_marimo_runtime() -> ( + None +): + """DefaultExecutor turns user exceptions into MarimoRuntimeException + with the user exception as __cause__. The teardown sees the wrapped + form, and the returned RunResult carries it as its exception.""" + from marimo._ast.cell import CellImpl + + log: list[str] = [] + a = _Recorder(log, "A") + + body_src = "raise ValueError('user bomb')" + + class _FakeCell: + cell_id = "0" + body = compile(body_src, "", "exec") + last_expr = compile("None", "", "eval") + + def is_coroutine(self) -> bool: + return False + + del CellImpl # silence unused-import + ev = Evaluator(executor=DefaultExecutor(), lifecycles=[a]) + result = await ev.evaluate(_FakeCell(), {}) # type: ignore[arg-type] + + assert isinstance(result.exception, MarimoRuntimeException) + assert isinstance(result.exception.__cause__, ValueError) + assert a.last_run_result is not None + # Teardown saw the wrapped exception, not the raw ValueError. + assert isinstance(a.last_run_result.exception, MarimoRuntimeException) + + +def _cause_traceback_filenames(exc: BaseException) -> list[str]: + cause = exc.__cause__ + assert cause is not None + tb = cause.__traceback__ + files: list[str] = [] + while tb is not None: + files.append(tb.tb_frame.f_code.co_filename) + tb = tb.tb_next + return files + + +def test_default_executor_strips_own_frame_from_cause_sync() -> None: + """`DefaultExecutor.execute_cell` must not leave its own frame on + the cause's `__traceback__` — user-facing tracebacks should begin + at user code (the compiled `` source).""" + + class _FakeCell: + cell_id = "0" + body = compile("raise ValueError('user bomb')", "", "exec") + last_expr = compile("None", "", "eval") + + with pytest.raises(MarimoRuntimeException) as exc_info: + DefaultExecutor().execute_cell(_FakeCell(), {}) # type: ignore[arg-type] + + files = _cause_traceback_filenames(exc_info.value) + assert files, "cause traceback unexpectedly empty" + assert not any("executor/executor.py" in f for f in files), files + assert files[0] == "" + + +async def test_default_executor_strips_own_frame_from_cause_async() -> None: + """Same as the sync variant, for `execute_cell_async`.""" + + class _FakeCell: + cell_id = "0" + body = compile("raise ValueError('user bomb')", "", "exec") + last_expr = compile("None", "", "eval") + + with pytest.raises(MarimoRuntimeException) as exc_info: + await DefaultExecutor().execute_cell_async(_FakeCell(), {}) # type: ignore[arg-type] + + files = _cause_traceback_filenames(exc_info.value) + assert files, "cause traceback unexpectedly empty" + assert not any("executor/executor.py" in f for f in files), files + assert files[0] == "" + + +async def test_teardown_runs_for_completed_setups_when_later_setup_raises() -> ( + None +): + log: list[str] = [] + a = _Recorder(log, "A") + b = _Recorder(log, "B", setup_raises=RuntimeError("setup-B raised")) + c = _Recorder(log, "C") # never reached + + ev = Evaluator( + executor=_StubExecutor(lambda *_: "ok"), + lifecycles=[a, b, c], + ) + result = await ev.evaluate(cell=None, glbls={}) + + assert isinstance(result.exception, RuntimeError) + assert str(result.exception) == "setup-B raised" + # A.setup ran (completed), B.setup ran and raised, C.setup did not + # run. Teardowns run only for lifecycles whose setup *completed* + # without raising — so only A. B is not teardowned because its + # state was never established. + assert log == [ + "setup:A", + "setup:B", + "teardown:A", + ] + + +async def test_teardown_wins_on_double_raise() -> None: + log: list[str] = [] + a = _Recorder(log, "A", teardown_raises=RuntimeError("teardown wins")) + + def body(cell: Any, glbls: dict[str, Any]) -> Any: + raise ValueError("body loses") + + ev = Evaluator(executor=_StubExecutor(body), lifecycles=[a]) + result = await ev.evaluate(cell=None, glbls={}) + + # Teardown exception replaces body exception in the final RunResult. + assert isinstance(result.exception, RuntimeError) + assert str(result.exception) == "teardown wins" + + +async def test_keyboard_interrupt_captured_into_run_result() -> None: + log: list[str] = [] + a = _Recorder(log, "A") + + def body(cell: Any, glbls: dict[str, Any]) -> Any: + raise KeyboardInterrupt + + ev = Evaluator(executor=_StubExecutor(body), lifecycles=[a]) + result = await ev.evaluate(cell=None, glbls={}) + + # Teardown ran (state still cleaned up) even though body raised + # BaseException, and the interrupt is captured in the RunResult + # rather than propagating out of evaluate(). + assert log == ["setup:A", "teardown:A"] + assert isinstance(result.exception, KeyboardInterrupt) + assert isinstance(a.last_run_result.exception, KeyboardInterrupt) + + +def test_strict_lifecycle_round_trip() -> None: + """Globals restored to pre-state after StrictLifecycle setup + + teardown.""" + from marimo._runtime.executor.lifecycles.strict import StrictLifecycle + + class _FakeCell: + cell_id = "c0" + refs: set[str] = set() + defs: set[str] = set() + + class _FakeGraph: + def get_transitive_references( + self, refs: set[str], predicate: Any + ) -> set[str]: + return set() + + lifecycle = StrictLifecycle(graph=_FakeGraph()) # type: ignore[arg-type] + glbls: dict[str, Any] = { + "x": 1, + "y": [1, 2, 3], + "__builtins__": __builtins__, + } + pre = {k: v for k, v in glbls.items()} + + skip = lifecycle.setup(_FakeCell(), glbls) # type: ignore[arg-type] + assert skip is None + + # During setup, glbls should be the sanitized scope (subset). + assert "x" not in glbls # No refs declared → x is not in scope. + + lifecycle.teardown(_FakeCell(), glbls, run_result=None) # type: ignore[arg-type] + + # Globals restored — same values for unchanged keys. + assert glbls["x"] == pre["x"] + assert glbls["y"] == pre["y"] + + +class _StrictGraph: + """`_FakeGraph` for `StrictLifecycle` setup-path tests. + + `transitive_refs` controls what `get_transitive_references` returns + so the test can drive `setup` past sanitization into the + error-construction branch. `defining_cells` maps refs to defining + cell IDs; refs absent from the map raise `KeyError` to exercise + the `unmangle_local` fallback. + """ + + def __init__( + self, + transitive_refs: set[str], + defining_cells: dict[str, list[str]] | None = None, + ) -> None: + self._transitive_refs = transitive_refs + self._defining_cells = defining_cells or {} + + def get_transitive_references( + self, refs: set[str], predicate: Any + ) -> set[str]: + return set(self._transitive_refs) + + def get_defining_cells(self, ref: str) -> list[str]: + return self._defining_cells[ref] + + +class _StrictCell: + def __init__(self, refs: set[str], defs: set[str] | None = None) -> None: + self.cell_id = "c0" + self.refs = refs + self.defs = defs or set() + + +def test_strict_setup_skip_on_undefined_ref() -> None: + """Unresolved ref → `Skip(result=RunResult(output=err, exception=err))` + where `err` is a `MarimoStrictExecutionError` with no blamed cell + (graph has no defining cell and the ref is not a private var).""" + from marimo._messaging.errors import MarimoStrictExecutionError + from marimo._runtime.executor.lifecycles.strict import StrictLifecycle + + lifecycle = StrictLifecycle( + graph=_StrictGraph(transitive_refs={"x"}) # type: ignore[arg-type] + ) + glbls: dict[str, Any] = {"__builtins__": {}} + + skip = lifecycle.setup(_StrictCell(refs={"x"}), glbls) # type: ignore[arg-type] + + assert skip is not None + assert skip.result is not None + err = skip.result.exception + assert isinstance(err, MarimoStrictExecutionError) + assert err.ref == "x" + assert err.blamed_cell is None + assert skip.result.output is err + + +def test_strict_setup_skip_on_ref_before_def() -> None: + """Ref appears in the cell's own `defs` → ref-before-def branch.""" + from marimo._messaging.errors import MarimoStrictExecutionError + from marimo._runtime.executor.lifecycles.strict import StrictLifecycle + + lifecycle = StrictLifecycle( + graph=_StrictGraph(transitive_refs={"x"}) # type: ignore[arg-type] + ) + glbls: dict[str, Any] = {"__builtins__": {}} + + skip = lifecycle.setup( + _StrictCell(refs={"x"}, defs={"x"}), # type: ignore[arg-type] + glbls, + ) + + assert skip is not None + assert skip.result is not None + err = skip.result.exception + assert isinstance(err, MarimoStrictExecutionError) + assert err.ref == "x" + assert err.blamed_cell is None + + +def test_strict_setup_skip_resolves_blamed_cell_via_graph() -> None: + """`get_defining_cells` returns the owning cell → blamed_cell.""" + from marimo._messaging.errors import MarimoStrictExecutionError + from marimo._runtime.executor.lifecycles.strict import StrictLifecycle + + lifecycle = StrictLifecycle( + graph=_StrictGraph( # type: ignore[arg-type] + transitive_refs={"x"}, + defining_cells={"x": ["other"]}, + ) + ) + glbls: dict[str, Any] = {"__builtins__": {}} + + skip = lifecycle.setup(_StrictCell(refs={"x"}), glbls) # type: ignore[arg-type] + + assert skip is not None + assert skip.result is not None + err = skip.result.exception + assert isinstance(err, MarimoStrictExecutionError) + assert err.blamed_cell == "other" + + +def test_strict_setup_skip_falls_back_to_private_var_owner() -> None: + """`KeyError` from the graph → `unmangle_local` resolves the + owning cell for mangled private vars.""" + from marimo._messaging.errors import MarimoStrictExecutionError + from marimo._runtime.executor.lifecycles.strict import StrictLifecycle + + # `_cell_ZZZ_priv` unmangles to (name="_priv", cell="ZZZ"). + private_ref = "_cell_ZZZ_priv" + lifecycle = StrictLifecycle( + graph=_StrictGraph( # type: ignore[arg-type] + transitive_refs={private_ref}, + ) + ) + glbls: dict[str, Any] = {"__builtins__": {}} + + skip = lifecycle.setup( + _StrictCell(refs={private_ref}), # type: ignore[arg-type] + glbls, + ) + + assert skip is not None + assert skip.result is not None + err = skip.result.exception + assert isinstance(err, MarimoStrictExecutionError) + assert err.blamed_cell == "ZZZ" + + +def test_strict_setup_skip_does_not_mutate_globals_or_stash_backup() -> None: + """The Skip early-return must happen before globals are cleared and + before the backup is stashed. `teardown` must then be a no-op.""" + from marimo._runtime.executor.lifecycles.strict import StrictLifecycle + + lifecycle = StrictLifecycle( + graph=_StrictGraph(transitive_refs={"x"}) # type: ignore[arg-type] + ) + glbls: dict[str, Any] = { + "preserve_me": 42, + "__builtins__": {}, + } + pre = dict(glbls) + + skip = lifecycle.setup(_StrictCell(refs={"x"}), glbls) # type: ignore[arg-type] + assert skip is not None + assert glbls == pre, "Skip path must not mutate globals" + assert lifecycle._backups == {}, "Skip path must not stash a backup" + + lifecycle.teardown(_StrictCell(refs={"x"}), glbls, skip.result) # type: ignore[arg-type] + assert glbls == pre, "teardown after Skip must be a no-op" + + +def test_execution_lifecycle_protocol_conformance() -> None: + """A Protocol-conforming class without inheriting works as a + lifecycle.""" + log: list[str] = [] + + class _MyLifecycle: + name = "mine" + + def setup(self, cell: Any, glbls: dict[str, Any]) -> Skip | None: + log.append("setup") + return None + + def teardown( + self, cell: Any, glbls: dict[str, Any], run_result: Any + ) -> None: + log.append("teardown") + + # Static type check via assignment to a ExecutionLifecycle-typed + # variable. If the Protocol is misshaped, mypy/pyright complains + # here, not at runtime. + lifecycle: ExecutionLifecycle = _MyLifecycle() + assert lifecycle.name == "mine" + + +# --- Surface 4: _cancel_on_sigint + evaluate_interruptible ------------------ + + +def _async_body(src: str) -> Any: + """Compile `src` with top-level-await support; returns a code object + whose `co_flags` carry `CO_COROUTINE` so `_is_coroutine` is True.""" + import ast + + return compile(src, "", "exec", flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT) + + +async def test_cancel_on_sigint_installs_and_restores_handler( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """`_cancel_on_sigint` swaps in its own handler on enter and + restores the previously-installed one on exit.""" + import signal + + from marimo._runtime.executor.evaluator import _cancel_on_sigint + + def prior(signum: int, frame: Any) -> None: + del signum, frame + + signal_calls: list[tuple[int, Any]] = [] + + def fake_signal(signum: int, handler: Any) -> Any: + signal_calls.append((signum, handler)) + return prior + + monkeypatch.setattr(signal, "signal", fake_signal) + monkeypatch.setattr(signal, "getsignal", lambda _signum: prior) + + fut: asyncio.Future[Any] = asyncio.Future() + with _cancel_on_sigint(fut): + # On enter: a new handler installed (not the prior). + assert signal_calls, "no signal.signal call recorded on enter" + assert signal_calls[0][0] == signal.SIGINT + assert signal_calls[0][1] is not prior + + # On exit: prior handler restored as the last call. + assert signal_calls[-1] == (signal.SIGINT, prior) + + +async def test_cancel_on_sigint_handler_cancels_future_and_chains_prior( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The installed handler must cancel the wrapped future and invoke + the previously-installed handler for its side effects.""" + import signal + + from marimo._runtime.executor.evaluator import _cancel_on_sigint + + prior_calls: list[tuple[int, Any]] = [] + + def prior(signum: int, frame: Any) -> None: + prior_calls.append((signum, frame)) + + captured: list[Any] = [] + + def fake_signal(signum: int, handler: Any) -> Any: + captured.append(handler) + return prior + + monkeypatch.setattr(signal, "signal", fake_signal) + monkeypatch.setattr(signal, "getsignal", lambda _signum: prior) + + fut: asyncio.Future[Any] = asyncio.Future() + with _cancel_on_sigint(fut): + marimo_handler = captured[0] + marimo_handler(signal.SIGINT, None) + # Cancellation propagates through done-callbacks asynchronously; + # yield to the loop so they fire. + await asyncio.sleep(0) + + assert fut.cancelled() + assert prior_calls == [(signal.SIGINT, None)] + + +async def test_cancel_on_sigint_swallows_marimo_interrupt_from_prior_handler( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Prior handler raising `MarimoInterrupt` must not escape — the + kernel's sync-mode raise is irrelevant for async cells, where the + halt comes from cancelling the future.""" + import signal + + from marimo._runtime.control_flow import MarimoInterrupt + from marimo._runtime.executor.evaluator import _cancel_on_sigint + + def prior(signum: int, frame: Any) -> None: + raise MarimoInterrupt + + captured: list[Any] = [] + + def fake_signal(signum: int, handler: Any) -> Any: + captured.append(handler) + return prior + + monkeypatch.setattr(signal, "signal", fake_signal) + monkeypatch.setattr(signal, "getsignal", lambda _signum: prior) + + fut: asyncio.Future[Any] = asyncio.Future() + with _cancel_on_sigint(fut): + marimo_handler = captured[0] + # No exception escapes — the wrapper catches MarimoInterrupt + # from the prior handler. + marimo_handler(signal.SIGINT, None) + await asyncio.sleep(0) + assert fut.cancelled() + + +async def test_executor_async_cancellation_propagates_unwrapped() -> None: + """`asyncio.CancelledError` must propagate unwrapped through + `DefaultExecutor.execute_cell_async` — wrapping it as + `MarimoRuntimeException` would mask the cancellation.""" + + class _AsyncCell: + cell_id = "0" + body = _async_body("import asyncio\nawait asyncio.sleep(100)") + last_expr = compile("None", "", "eval") + + def is_coroutine(self) -> bool: + return True + + task = asyncio.create_task( + DefaultExecutor().execute_cell_async(_AsyncCell(), {}) # type: ignore[arg-type] + ) + # Yield so the task enters the awaited sleep before we cancel. + await asyncio.sleep(0) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + +async def test_evaluate_interruptible_no_op_for_sync_cell() -> None: + """Sync cells: `evaluate_interruptible` returns the same shape as a + direct `evaluate()` call. The SIGINT-handler wrap is for async only.""" + + class _SyncCell: + cell_id = "0" + body = compile("x = 1", "", "exec") + last_expr = compile("x", "", "eval") + + def is_coroutine(self) -> bool: + return False + + ev = Evaluator(executor=DefaultExecutor(), lifecycles=[]) + + sync_glbls: dict[str, Any] = {} + interruptible_glbls: dict[str, Any] = {} + + direct = await ev.evaluate(_SyncCell(), sync_glbls) # type: ignore[arg-type] + interruptible = await ev.evaluate_interruptible( + _SyncCell(), # type: ignore[arg-type] + interruptible_glbls, + ) + + assert direct.output == interruptible.output == 1 + assert direct.exception is None + assert interruptible.exception is None + assert direct.accumulated_output == interruptible.accumulated_output diff --git a/tests/_runtime/test_scheduler.py b/tests/_runtime/test_scheduler.py new file mode 100644 index 00000000000..58a0a517ea6 --- /dev/null +++ b/tests/_runtime/test_scheduler.py @@ -0,0 +1,82 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Queue + cancellation invariants for SequentialScheduler.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import MagicMock + +from marimo._runtime.runner.scheduler import SequentialScheduler +from marimo._types.ids import CellId_t + +if TYPE_CHECKING: + import pytest + + +def _empty_graph() -> MagicMock: + """A graph whose transitive_closure returns just the input cell.""" + g = MagicMock() + g.cells = {} + return g + + +def test_pending_and_pop_cell_fifo() -> None: + cells = [CellId_t("a"), CellId_t("b"), CellId_t("c")] + sched = SequentialScheduler(cells, graph=_empty_graph()) + + assert sched.pending() is True + assert sched.pop_cell() == "a" + assert sched.pop_cell() == "b" + assert sched.pop_cell() == "c" + assert sched.pending() is False + + +def test_interrupted_blocks_pending() -> None: + sched = SequentialScheduler([CellId_t("a")], graph=_empty_graph()) + + assert sched.pending() is True + sched.interrupted = True + assert sched.pending() is False + + +def test_cancel_marks_cancelled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Mock graph: transitive_closure returns just the cell itself, no + # descendants. Cell registered in graph.cells so set_run_result_status + # has a target. + g = MagicMock() + cid = CellId_t("a") + cell_mock = MagicMock() + g.cells = {cid: cell_mock} + + def fake_closure(graph: object, roots: set[CellId_t]) -> set[CellId_t]: + del graph + return set(roots) + + monkeypatch.setattr( + "marimo._runtime.dataflow.transitive_closure", fake_closure + ) + sched = SequentialScheduler([cid], graph=g) + assert sched.cancelled(cid) is False + sched.cancel(cid) + assert sched.cancelled(cid) is True + cell_mock.set_run_result_status.assert_called_with("cancelled") + + +def test_batch_yields_singletons() -> None: + sched = SequentialScheduler([], graph=_empty_graph()) + cells = [CellId_t("a"), CellId_t("b"), CellId_t("c")] + batches = list(sched.batch(cells)) + assert batches == [["a"], ["b"], ["c"]] + + +def test_batch_respects_interrupt() -> None: + sched = SequentialScheduler([], graph=_empty_graph()) + cells = [CellId_t("a"), CellId_t("b"), CellId_t("c")] + iterator = sched.batch(cells) + assert next(iterator) == ["a"] + sched.interrupted = True + # Generator stops once interrupted is set. + remaining = list(iterator) + assert remaining == [] diff --git a/tests/_runtime/test_unwrap_user_exception.py b/tests/_runtime/test_unwrap_user_exception.py new file mode 100644 index 00000000000..ad68a2fc76b --- /dev/null +++ b/tests/_runtime/test_unwrap_user_exception.py @@ -0,0 +1,85 @@ +# Copyright 2026 Marimo. All rights reserved. +"""Unit tests for `marimo._runtime.exceptions.unwrap_user_exception`.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +from marimo._runtime.exceptions import ( + MarimoMissingRefError, + MarimoRuntimeException, + unwrap_user_exception, +) + + +def _wrap(cause: BaseException) -> MarimoRuntimeException: + """Build a `MarimoRuntimeException` with `__cause__` set.""" + try: + raise MarimoRuntimeException from cause + except MarimoRuntimeException as exc: + return exc + + +def _graph(definitions: set[str]) -> Any: + """Minimal graph stub with only the attribute `unwrap` reads.""" + return SimpleNamespace(definitions=definitions) + + +def test_unwrap_no_graph_returns_raw_cause() -> None: + cause = ValueError("boom") + wrapped = _wrap(cause) + + assert unwrap_user_exception(wrapped) is cause + + +def test_unwrap_nameerror_without_graph_unchanged() -> None: + """No graph → upgrade never fires, even for NameError.""" + cause = NameError("name 'x' is not defined") + cause.name = "x" # set explicitly; constructor doesn't. + wrapped = _wrap(cause) + + assert unwrap_user_exception(wrapped) is cause + + +def test_unwrap_nameerror_with_graph_upgrades_when_in_definitions() -> None: + cause = NameError("name 'x' is not defined") + cause.name = "x" + wrapped = _wrap(cause) + + unwrapped = unwrap_user_exception(wrapped, graph=_graph({"x"})) + + assert isinstance(unwrapped, MarimoMissingRefError) + assert unwrapped.ref == "x" + assert unwrapped.name_error is cause + + +def test_unwrap_nameerror_with_graph_passthrough_when_not_in_definitions() -> ( + None +): + """`.name` is set but the graph doesn't define it — no upgrade.""" + cause = NameError("name 'x' is not defined") + cause.name = "x" + wrapped = _wrap(cause) + + assert unwrap_user_exception(wrapped, graph=_graph(set())) is cause + + +def test_unwrap_nameerror_with_none_name_returns_raw() -> None: + """`NameError.name is None` (the constructor default) → upgrade + short-circuits via the `if name and …` guard.""" + cause = NameError("name 'x' is not defined") + # Don't set `.name` — leave it as the constructor's default + # (None on most CPython versions). The guard must not upgrade. + assert getattr(cause, "name", None) is None + wrapped = _wrap(cause) + + # Even with `x` in graph.definitions, the guard prevents upgrade. + assert unwrap_user_exception(wrapped, graph=_graph({"x"})) is cause + + +def test_unwrap_no_cause_returns_none() -> None: + """`MarimoRuntimeException` raised without `from …` has no cause.""" + wrapped = MarimoRuntimeException() + + assert unwrap_user_exception(wrapped) is None diff --git a/tests/_server/test_scratchpad_integration.py b/tests/_server/test_scratchpad_integration.py index c118976c3ee..da5a1d716b8 100644 --- a/tests/_server/test_scratchpad_integration.py +++ b/tests/_server/test_scratchpad_integration.py @@ -672,7 +672,7 @@ def test_ctx_create_cell_multiply_defined(session: _Session) -> None: assert lines == snapshot( [ "event: stderr", - 'data: {"data": "Traceback (most recent call last):\\n File \\"/marimo/_runtime/executor.py\\", line N, in execute_cell_async\\n await eval(cell.body, glbls)\\n File \\"\\", line 2, in \\n async with cm.get_context() as ctx:\\n File \\"/marimo/_code_mode/_context.py\\", line N, in __aexit__\\n self._dry_run_compile(ops)\\n File \\"/marimo/_code_mode/_context.py\\", line N, in _dry_run_compile\\n raise RuntimeError(\\nRuntimeError: Multiply-defined names:\\n - \'x\' is already defined in cell \'cell_a\' (cell_a)\\n\\nTo skip validation, use: async with cm.get_context(skip_validation=True) as ctx\\n"}', + 'data: {"data": "Traceback (most recent call last):\\n File \\"\\", line 2, in \\n async with cm.get_context() as ctx:\\n File \\"/marimo/_code_mode/_context.py\\", line N, in __aexit__\\n self._dry_run_compile(ops)\\n File \\"/marimo/_code_mode/_context.py\\", line N, in _dry_run_compile\\n raise RuntimeError(\\nRuntimeError: Multiply-defined names:\\n - \'x\' is already defined in cell \'cell_a\' (cell_a)\\n\\nTo skip validation, use: async with cm.get_context(skip_validation=True) as ctx\\n"}', "", "event: done", 'data: {"success": false, "output": {"mimetype": "text/plain", "data": ""}}', diff --git a/tests/test_entrypoints.py b/tests/test_entrypoints.py index c39a29748a2..135301cc57d 100644 --- a/tests/test_entrypoints.py +++ b/tests/test_entrypoints.py @@ -1,12 +1,19 @@ +from __future__ import annotations + import os -from typing import cast +from typing import TYPE_CHECKING, Any, cast from unittest.mock import MagicMock, patch import pytest from marimo._entrypoints.ids import KnownEntryPoint from marimo._entrypoints.registry import EntryPointRegistry, get_entry_points -from marimo._runtime.executor import ExecutionConfig, Executor, get_executor +from marimo._runtime.executor import Executor + +if TYPE_CHECKING: + from collections.abc import Callable + + from marimo._ast.cell import CellImpl class TestEntryPointRegistry: @@ -166,35 +173,77 @@ def test_get_all_with_entry_points( assert set(result) == {"value1", "ep_value1", "ep_value2"} -class CustomExecutor(Executor): +class CustomExecutor: + """Protocol-conforming Executor (no ABC inheritance).""" + + name = "custom" + def execute_cell( self, - cell: str, - glbls: dict[str, str], - graph: str, - ) -> str: - return f"Executed {cell} with {glbls} in {graph}" + cell: CellImpl, + glbls: dict[str, Any], + ) -> Any: + return f"Executed {cell} with {glbls}" async def execute_cell_async( self, - cell: str, - glbls: dict[str, str], - graph: str, - ) -> str: - return f"Executed {cell} with {glbls} in {graph}" + cell: CellImpl, + glbls: dict[str, Any], + ) -> Any: + return f"Executed {cell} with {glbls}" -class TestExecutorEntryPoint: - @pytest.fixture - def registry(self) -> EntryPointRegistry[Executor]: - reg = EntryPointRegistry[Executor]("marimo.cell.executor") - reg.register("custom", CustomExecutor) - return reg +def _custom_executor_factory() -> Executor: + return CustomExecutor() - def test_get_entry_points_modern( - self, registry: EntryPointRegistry[Executor] - ) -> None: - executor = get_executor( - ExecutionConfig(is_strict=False), registry=registry + +class TestExecutorEntryPoint: + def test_factory_registers_and_resolves(self) -> None: + # Registry holds factories (Callable[[], Executor]); the kernel + # calls the factory once to get an instance. + reg: EntryPointRegistry[Callable[[], Executor]] = EntryPointRegistry( + "marimo.cell.executor" ) + reg.register("custom", _custom_executor_factory) + + factory = reg.get("custom") + executor = factory() assert isinstance(executor, CustomExecutor) + assert executor.execute_cell("c", {"x": "1"}) == ( # type: ignore[arg-type] + "Executed c with {'x': '1'}" + ) + + def test_resolve_executor_only_loads_first_factory(self) -> None: + """`resolve_executor` must not import factories beyond the first. + + A broken or slow third-party plugin can't take down the kernel + if it never gets loaded. + """ + from marimo._runtime.executor.evaluator import ( + _EXECUTOR_REGISTRY, + resolve_executor, + ) + + loaded: list[str] = [] + + def working_factory() -> Executor: + loaded.append("working") + return CustomExecutor() + + def broken_factory() -> Executor: + loaded.append("broken") + raise RuntimeError("third-party plugin is broken") + + # Restore the registry's plugins on exit so we don't leak + # registrations into other tests. + before = dict(_EXECUTOR_REGISTRY._plugins) + _EXECUTOR_REGISTRY._plugins.clear() + try: + _EXECUTOR_REGISTRY.register("aaa-working", working_factory) + _EXECUTOR_REGISTRY.register("zzz-broken", broken_factory) + executor = resolve_executor() + assert isinstance(executor, CustomExecutor) + assert loaded == ["working"] + finally: + _EXECUTOR_REGISTRY._plugins.clear() + _EXECUTOR_REGISTRY._plugins.update(before)