feat: top level imports

marimo/__init__.py

@@ -48,6 +48,7 @@
     "icon",
     "iframe",
     "image",
+    "import_guard",
     "latex",
     "lazy",
     "left",
@@ -125,7 +126,7 @@
     redirect_stderr,
     redirect_stdout,
 )
-from marimo._runtime.context.utils import running_in_notebook
+from marimo._runtime.context.utils import import_guard, running_in_notebook
 from marimo._runtime.control_flow import MarimoStopError, stop
 from marimo._runtime.runtime import (
     app_meta,

marimo/_ast/app.py

@@ -21,9 +21,12 @@
     TypeVar,
     Union,
     cast,
+    overload,
 )
 from uuid import uuid4
 
+from typing_extensions import ParamSpec, TypeAlias
+
 from marimo import _loggers
 from marimo._ast.cell import Cell, CellConfig, CellId_t, CellImpl
 from marimo._ast.cell_manager import CellManager
@@ -58,7 +61,9 @@
     from marimo._runtime.context.types import ExecutionContext
 
 
-Fn = TypeVar("Fn", bound=Callable[..., Any])
+P = ParamSpec("P")
+R = TypeVar("R")
+Fn: TypeAlias = Callable[P, R]
 LOGGER = _loggers.marimo_logger()
 
 
@@ -266,13 +271,13 @@ def clone(self) -> App:
 
     def cell(
         self,
-        func: Fn | None = None,
+        func: Fn[P, R] | None = None,
         *,
         column: Optional[int] = None,
         disabled: bool = False,
         hide_code: bool = False,
         **kwargs: Any,
-    ) -> Cell | Callable[[Fn], Cell]:
+    ) -> Cell | Callable[[Fn[P, R]], Cell]:
         """A decorator to add a cell to the app.
 
         This decorator can be called with or without parentheses. Each of the
@@ -302,21 +307,29 @@ def __(mo):
         del kwargs
 
         return cast(
-            Union[Cell, Callable[[Fn], Cell]],
+            Union[Cell, Callable[[Fn[P, R]], Cell]],
             self._cell_manager.cell_decorator(
                 func, column, disabled, hide_code, app=InternalApp(self)
             ),
         )
 
+    # Overloads are required to preserve the wrapped function's signature.
+    # mypy is not smart enough to carry transitive typing in this case.
+    @overload
+    def function(self, func: Fn[P, R]) -> Fn[P, R]: ...
+
+    @overload
+    def function(self, **kwargs: Any) -> Callable[[Fn[P, R]], Fn[P, R]]: ...
+
     def function(
         self,
-        func: Fn | None = None,
+        func: Fn[P, R] | None = None,
         *,
         column: Optional[int] = None,
         disabled: bool = False,
         hide_code: bool = False,
         **kwargs: Any,
-    ) -> Fn | Callable[[Fn], Fn]:
+    ) -> Fn[P, R] | Callable[[Fn[P, R]], Fn[P, R]]:
         """A decorator to wrap a callable function into a cell in the app.
 
         This decorator can be called with or without parentheses. Each of the
@@ -348,7 +361,7 @@ def multiply(a: int, b: int) -> int:
         del kwargs
 
         return cast(
-            Union[Fn, Callable[[Fn], Fn]],
+            Union[Fn[P, R], Callable[[Fn[P, R]], Fn[P, R]]],
             self._cell_manager.cell_decorator(
                 func,
                 column,

marimo/_ast/cell.py

@@ -289,8 +289,9 @@ def namespace_to_variable(self, namespace: str) -> Name | None:
     def is_coroutine(self) -> bool:
         return _is_coroutine(self.body) or _is_coroutine(self.last_expr)
 
-    @property
-    def is_toplevel_acceptable(self) -> bool:
+    def is_toplevel_acceptable(
+        self, allowed_refs: Optional[set[Name]] = None
+    ) -> bool:
         # Check no defs aside from the single function
         if len(self.defs) != 1:
             return False
@@ -316,13 +317,16 @@ def is_toplevel_acceptable(self) -> bool:
             return False
 
         # No required_refs are allowed for now
-        # TODO: Allow imports and other toplevel functions
         refs = set().union(
             *[v.required_refs for v in self.variable_data[name]]
         )
-        refs -= set(globals()["__builtins__"].keys())
+        # NOTE: Builtins are allowed, but should be passed in under
+        # allowed_refs. Defers to allowed_refs because shadowed builtins
+        # are accounted for.
+        if allowed_refs is None:
+            allowed_refs = set(globals()["__builtins__"].keys())
         # Allow recursion
-        refs -= {name}
+        refs -= {name} | allowed_refs
         if refs:
             return False
 

marimo/_ast/cell_manager.py

@@ -7,13 +7,14 @@
 import string
 from typing import (
     TYPE_CHECKING,
-    Any,
     Callable,
     Iterable,
     Optional,
     TypeVar,
 )
 
+from typing_extensions import ParamSpec, TypeAlias
+
 from marimo._ast.cell import Cell, CellConfig, CellId_t
 from marimo._ast.compiler import cell_factory, toplevel_cell_factory
 from marimo._ast.models import CellData
@@ -23,7 +24,9 @@
 if TYPE_CHECKING:
     from marimo._ast.app import InternalApp
 
-    Fn = TypeVar("Fn", bound=Callable[..., Any])
+P = ParamSpec("P")
+R = TypeVar("R")
+Fn: TypeAlias = Callable[P, R]
 
 
 class CellManager:
@@ -74,20 +77,20 @@ def create_cell_id(self) -> CellId_t:
     # TODO: maybe remove this, it is leaky
     def cell_decorator(
         self,
-        func: Fn | None,
+        func: Fn[P, R] | None,
         column: Optional[int],
         disabled: bool,
         hide_code: bool,
         app: InternalApp | None = None,
         *,
         top_level: bool = False,
-    ) -> Cell | Fn | Callable[[Fn], Cell | Fn]:
+    ) -> Cell | Fn[P, R] | Callable[[Fn[P, R]], Cell | Fn[P, R]]:
         """Create a cell decorator for marimo notebook cells."""
         cell_config = CellConfig(
             column=column, disabled=disabled, hide_code=hide_code
         )
 
-        def _register(func: Fn) -> Cell | Fn:
+        def _register(func: Fn[P, R]) -> Cell | Fn[P, R]:
             # Use PYTEST_VERSION here, opposed to PYTEST_CURRENT_TEST, in
             # order to allow execution during test collection.
             is_top_level_pytest = (
@@ -111,16 +114,15 @@ def _register(func: Fn) -> Cell | Fn:
 
             # Manually set the signature for pytest.
             if is_top_level_pytest:
-                func = wrap_fn_for_pytest(func, cell)  # type: ignore
-            # NB. in place metadata update.
-            functools.wraps(func)(cell)
+                # NB. in place metadata update.
+                functools.wraps(wrap_fn_for_pytest(func, cell))(cell)
             return cell
 
         if func is None:
             # If the decorator was used with parentheses, func will be None,
             # and we return a decorator that takes the decorated function as an
             # argument
-            def decorator(func: Fn) -> Cell | Fn:
+            def decorator(func: Fn[P, R]) -> Cell | Fn[P, R]:
                 return _register(func)
 
             return decorator