feat: Add markdown to the cell level

marimo/_ast/cell.py

@@ -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)

marimo/_ast/compiler.py

@@ -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,
@@ -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),
@@ -310,6 +373,7 @@ def compile_cell(
         body=body,
         last_expr=last_expr,
         cell_id=cell_id,
+        markdown=maybe_md,
         _test=is_test,
     )
 

marimo/_runtime/runtime.py

@@ -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
+
+        # Since markdown cell, render and broadcast output
+        # Remove cell from outstanding requests
+        from marimo._output.md import md
+
+        # Remove markdown cells from uninstantiated requests
+        for cell_id, content in markdown_cells.items():
+            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]
 
     def load_dotenv(self) -> None:
         dotenvs = self.user_config["runtime"].get("dotenv", [])

marimo/_server/export/utils.py

@@ -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:
@@ -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]]:
@@ -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"