Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions marimo/_ast/cell.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,9 @@ class CellImpl:
# unique id
cell_id: CellId_t

# Markdown content of the cell if it exists
markdown: Optional[str] = None

# Mutable fields
# explicit configuration of cell
config: CellConfig = dataclasses.field(default_factory=CellConfig)
Expand Down
64 changes: 64 additions & 0 deletions marimo/_ast/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,67 @@ def fix_source_position(node: Any, source_position: SourcePosition) -> Any:
return node


def const_string(args: list[ast.stmt]) -> str:
(inner,) = args
if hasattr(inner, "values"):
(inner,) = inner.values
return f"{inner.value}" # type: ignore[attr-defined]


def const_or_id(args: ast.stmt) -> str:
if hasattr(args, "value"):
return f"{args.value}" # type: ignore[attr-defined]
return f"{args.id}" # type: ignore[attr-defined]


def _extract_markdown(tree: ast.Module) -> Optional[str]:
# Attribute Error handled by the outer try/except block.
# Wish there was a more compact to ignore ignore[attr-defined] for all.
try:
(body,) = tree.body
if body.value.func.attr == "md": # type: ignore[attr-defined, union-attr]
value = body.value # type: ignore[attr-defined, union-attr]
else:
return None
assert value.func.value.id == "mo"
if not value.args: # Handle mo.md() with no arguments
return None
md_lines = const_string(value.args).split("\n")
except (AssertionError, AttributeError, ValueError):
# No reason to explicitly catch exceptions if we can't parse out
# markdown. Just handle it as a code block.
return None

# Dedent behavior is a little different that in marimo js, so handle
# accordingly.
md_lines = [line.rstrip() for line in md_lines]
md = (
textwrap.dedent(md_lines[0])
+ "\n"
+ textwrap.dedent("\n".join(md_lines[1:]))
)
md = md.strip()
return md


def extract_markdown(code: str) -> Optional[str]:
code = code.strip()
count = 0
# Early quitting for markdown extraction.
for line in code.strip().split("\n"):
if line.startswith("mo.md("):
count += 1
if count > 1:
return None
if count == 0:
return None

try:
return _extract_markdown(ast.parse(code))
except SyntaxError:
return None


def compile_cell(
code: str,
cell_id: CellId_t,
Expand Down Expand Up @@ -291,6 +352,8 @@ def compile_cell(
if previous_import_data == import_data:
imported_defs.add(import_data.definition)

maybe_md = _extract_markdown(original_module)

return CellImpl(
# keyed by original (user) code, for cache lookups
key=code_key(code),
Expand All @@ -310,6 +373,7 @@ def compile_cell(
body=body,
last_expr=last_expr,
cell_id=cell_id,
markdown=maybe_md,
_test=is_test,
)

Expand Down
87 changes: 80 additions & 7 deletions marimo/_runtime/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -2009,22 +2009,95 @@ async def instantiate(self, request: CreationRequest) -> None:
if self.graph.cells:
del request
LOGGER.debug("App already instantiated.")
elif request.auto_run:
return

# Handle markdown cells specially during kernel-ready initialization
execution_requests = {
er.cell_id: er for er in request.execution_requests
}
self._handle_markdown_cells_on_instantiate(execution_requests)

if request.auto_run:
self.reset_ui_initializers()
for (
object_id,
initial_value,
) in request.set_ui_element_value_request.ids_and_values:
self.ui_initializers[object_id] = initial_value

await self.run(request.execution_requests)
await self.run(list(execution_requests.values()))
self.reset_ui_initializers()
else:
self._uninstantiated_execution_requests = {
er.cell_id: er for er in request.execution_requests
}
for cid in self._uninstantiated_execution_requests:
CellOp.broadcast_stale(cell_id=cid, stale=True)
self._uninstantiated_execution_requests = execution_requests
for cell_id in self._uninstantiated_execution_requests.keys():
CellOp.broadcast_stale(cell_id=cell_id, stale=True)

def _handle_markdown_cells_on_instantiate(
self, execution_requests: dict[CellId_t, ExecutionRequest]
) -> None:
"""Handle markdown cells during kernel-ready initialization.

For cells that contain only markdown (mo.md calls), this method:
1. Compiles the cells to extract markdown content
2. Renders the markdown to HTML
3. Broadcasts the rendered output immediately
4. Marks the cells as completed (not stale)
5. Removes them from uninstantiated requests

NOTE: If 'mo' is not available in the graph definitions, all cells are
marked as stale. Regular cells are marked as stale as usual.
"""
# If 'mo' is not available in the graph, mark all cells as stale
markdown_cells: dict[CellId_t, str] = {}
exports_mo = False
for cid, er in execution_requests.items():
# Check if cell already exists in graph (to avoid recompilation)
cell = self.graph.cells.get(cid)
error = None

# If cell doesn't exist in graph, try to compile it
if cell is None:
# TODO: Don't bother compiling whole cell.
# However, since we still need to extract defs
# for mo / marimo, this is OK for now.
cell, error = self._try_compiling_cell(cid, er.code, [])

if cell is None or error is not None:
continue

# Check if this is a markdown cell
if cell.markdown is not None:
# Remove from uninstantiated requests since it's effectively "run"
markdown_cells[cid] = cell.markdown
else:
# Regular cell - mark as stale
exports_mo |= "mo" in cell.defs

# Handle as default if no cells export 'mo'
if not exports_mo:
return

# Remove markdown cells from uninstantiated requests
for cell_id, content in markdown_cells.items():
# Since markdown cell, render and broadcast output
# Remove cell from outstanding requests
from marimo._output.md import md

html_obj = md(content)
mimetype, html_content = html_obj._mime_()

# Broadcast the markdown output
CellOp.broadcast_output(
channel=CellChannel.OUTPUT,
mimetype=mimetype,
data=html_content,
cell_id=cell_id,
status="idle",
)

# Mark the cell as not stale (already "run")
CellOp.broadcast_stale(cell_id=cell_id, stale=False)
del execution_requests[cell_id]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

any way to break some of this up outside of the Runner? its gotten quite large. i know it doesn't really belong in hooks, but maybe a function in the same file

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good point, but we still call some functions on the object


def load_dotenv(self) -> None:
dotenvs = self.user_config["runtime"].get("dotenv", [])
Expand Down
46 changes: 3 additions & 43 deletions marimo/_server/export/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@
import ast
import os
import re
from textwrap import dedent
from typing import Optional, Union

from marimo._ast.cell import Cell, CellImpl
from marimo._ast.compiler import const_or_id, extract_markdown


def format_filename_title(filename: str) -> str:
Expand All @@ -31,54 +31,14 @@ def get_download_filename(filename: Optional[str], extension: str) -> str:
return f"{os.path.splitext(basename)[0]}.{extension}"


def _const_string(args: list[ast.stmt]) -> str:
(inner,) = args
if hasattr(inner, "values"):
(inner,) = inner.values
return f"{inner.value}" # type: ignore[attr-defined]


def _const_or_id(args: ast.stmt) -> str:
if hasattr(args, "value"):
return f"{args.value}" # type: ignore[attr-defined]
return f"{args.id}" # type: ignore[attr-defined]


def get_markdown_from_cell(
cell: Union[CellImpl, Cell], code: str
) -> Optional[str]:
"""Attempt to extract markdown from a cell, or return None"""

if not (cell.refs == {"mo"} and not cell.defs):
return None
markdown_lines = [
line for line in code.strip().split("\n") if line.startswith("mo.md(")
]
if len(markdown_lines) > 1:
return None

code = code.strip()
# Attribute Error handled by the outer try/except block.
# Wish there was a more compact to ignore ignore[attr-defined] for all.
try:
(body,) = ast.parse(code).body
if body.value.func.attr == "md": # type: ignore[attr-defined]
value = body.value # type: ignore[attr-defined]
else:
return None
assert value.func.value.id == "mo"
md_lines = _const_string(value.args).split("\n")
except (AssertionError, AttributeError, ValueError, SyntaxError):
# No reason to explicitly catch exceptions if we can't parse out
# markdown. Just handle it as a code block.
return None

# Dedent behavior is a little different that in marimo js, so handle
# accordingly.
md_lines = [line.rstrip() for line in md_lines]
md = dedent(md_lines[0]) + "\n" + dedent("\n".join(md_lines[1:]))
md = md.strip()
return md
return extract_markdown(code)


def get_sql_options_from_cell(code: str) -> Optional[dict[str, str]]:
Expand All @@ -96,7 +56,7 @@ def get_sql_options_from_cell(code: str) -> Optional[dict[str, str]]:
return None
if value.keywords:
for keyword in value.keywords: # type: ignore[attr-defined]
options[keyword.arg] = _const_or_id(keyword.value) # type: ignore[attr-defined]
options[keyword.arg] = const_or_id(keyword.value) # type: ignore[attr-defined]
output = options.pop("output", "True").lower()
if output == "false":
options["hide_output"] = "True"
Expand Down
Loading
Loading