feat: top level imports (#3779)

## 📝 Summary

Followup to #3755 for #2293 allowing for "top level imports"

For completion of #2293, I thin UI changes and needed for enabling this
behavior. Notably:

   - Indicate when in function mode (maybe top level import too)
   - Provide hints when pushed out of function mode
   - Maybe allow the user to opt out of function mode?
 
\+ docs

This also increases security risk since code is run outside of runtime.
This was always possible, but now marimo can save in a format that could
skip the marimo runtime all together on restart.

There are opportunities here. marimo could lean into this, and leverage
external code running as a chance to hook in (almost a plugin system for
free)

But also issues, since a missing dep could stop the notebook from
running at all (goes against the "batteries included" ethos). This can
be mitigated with static analysis over just an import (markdown does
this for instance), or marimo can re-serialize the notebook in the
"safe" form, if it comes across issues in import.

## 🔍 Description of Changes

Includes a bit of a refactor to codegen since there were a fair amount
of changes.
Allows top level imports of "import only" cells. The contents are pasted
at the top of the file, with a bit of care not to break header
extraction.

```python
# Normal headers are retained
# Use a notice to denote where generated imports start
# Notice maybe needs some copy edit

# 👋 This file was generated by marimo. You can edit it, and tweak
# things- just be conscious that some changes may be overwritten if opened in
# the editor. For instance top level imports are derived from a cell, and not
# the top of the script. This notice signifies the beginning of the generated
# import section.

# Could also make this app.imports? But maybe increasing surface area for no reason
import numpy
# Note, import cells intentionally do not have a `return`
# for static analysis feature below

import marimo


__generated_with = "0.11.2"
app = marimo.App(_toplevel_fn=True)


@app.cell
def import_cell():
    # Could also make this app.imports? But maybe increasing surface area for no reason
    import numpy
    # Note, import cells intentionally do not have a `return`
    # for static analysis feature below

```

Top level refs (this includes `@app.function`s) are ignored in the
signatures. E.g.

```python
import marimo as mo

# ...

@app.cell
def md_cell():
    mo.md("Hi")
    return 
```

Since I was also in there, I added static analysis to ignore returning
dangling defs.

```python
@app.cell
def cell_with_dangling_def():
    a = 1
    b = 2
    return (a,) # No longer returns b since it's not used anywhere. Allowing for linters like ruff to complain.

@app.cell
def ref_cell(a):
    a + 1
    return 
```

LMK if too far reaching and we can break it up/ refactor. A bit more
opinionated than the last PR

Test border more on being more smoke tests than unit tests, but hit the
key issues I was worried about. I can break them down more granularly if
needed. Also LMK if you can think of some more edgecases.

## 📜 Reviewers

<!--
Tag potential reviewers from the community or maintainers who might be
interested in reviewing this pull request.

Your PR will be reviewed more quickly if you can figure out the right
person to tag with @ -->

@akshayka OR @mscolnick

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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