Extract and lift inline script metadata in marimo convert

Author: manzt

marimo/_convert/ipynb.py

@@ -3,9 +3,10 @@
 
 import ast
 import json
+import re
 import sys
 from collections import defaultdict
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Union
 
 from marimo._ast.compiler import compile_cell
 from marimo._ast.transformers import NameTransformer
@@ -603,6 +604,21 @@ def transform_strip_whitespace(sources: List[str]) -> List[str]:
     return [source.strip() for source in sources]
 
 
+def extract_inline_meta(script: str) -> tuple[str | None, str]:
+    """
+    Extract PEP 723 metadata from a Python source.
+
+    Returns a tuple of the metadata comment and the remaining script.
+    """
+    if match := re.search(
+        r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s(?P<content>(^#(| .*)$\s)+)^# ///$",
+        script,
+    ):
+        meta_comment = match.group(0)
+        return meta_comment, script.replace(meta_comment, "").strip()
+    return None, script
+
+
 def _transform_sources(
     sources: list[str], metadata: list[dict[str, Any]]
 ) -> list[str]:
@@ -627,6 +643,7 @@ def convert_from_ipynb(raw_notebook: str) -> str:
     notebook = json.loads(raw_notebook)
     sources: List[str] = []
     metadata: List[Dict[str, Any]] = []
+    inline_meta: Union[str, None] = None
 
     for cell in notebook["cells"]:
         source = (
@@ -636,7 +653,14 @@ def convert_from_ipynb(raw_notebook: str) -> str:
         )
         if cell["cell_type"] == "markdown":
             source = markdown_to_marimo(source)
-        sources.append(source)
-        metadata.append(cell.get("metadata", {}))
+        elif inline_meta is None:
+            # Eagerly find PEP 723 metadata, first match wins
+            inline_meta, source = extract_inline_meta(source)
+
+        if source:
+            sources.append(source)
+            metadata.append(cell.get("metadata", {}))
 
-    return generate_from_sources(_transform_sources(sources, metadata))
+    return generate_from_sources(
+        _transform_sources(sources, metadata), header_comments=inline_meta
+    )

marimo/_convert/utils.py

@@ -23,7 +23,9 @@ def markdown_to_marimo(source: str) -> str:
 
 
 def generate_from_sources(
-    sources: list[str], config: Optional[_AppConfig] = None
+    sources: list[str],
+    config: Optional[_AppConfig] = None,
+    header_comments: Optional[str] = None,
 ) -> str:
     """
     Given a list of Python source code,
@@ -34,4 +36,5 @@ def generate_from_sources(
         ["__" for _ in sources],
         [CellConfig() for _ in range(len(sources))],
         config=config,
+        header_comments=header_comments,
     )