refactor: isolate markdown import dialects

Author: peter-gy

marimo/_convert/markdown/flavor/__init__.py

@@ -9,8 +9,12 @@
 from marimo._convert.markdown.flavor.base import (
     MarkdownFlavor,
     MarkdownFlavorName,
+    MarkdownImportDialect,
+)
+from marimo._convert.markdown.flavor.mystmd import (
+    MystmdMarkdownFlavor,
+    _MystmdMarkdownImportDialect,
 )
-from marimo._convert.markdown.flavor.mystmd import MystmdMarkdownFlavor
 from marimo._convert.markdown.flavor.pymdown import PymdownMarkdownFlavor
 from marimo._convert.markdown.flavor.qmd import QmdMarkdownFlavor
 
@@ -20,6 +24,7 @@
 _PYMDOWN_MARKDOWN = PymdownMarkdownFlavor()
 _QMD_MARKDOWN = QmdMarkdownFlavor()
 _MYSTMD_MARKDOWN = MystmdMarkdownFlavor()
+_MYSTMD_MARKDOWN_IMPORT = _MystmdMarkdownImportDialect()
 _MARKDOWN_FLAVORS: Mapping[MarkdownFlavorName, MarkdownFlavor] = (
     MappingProxyType(
         {
@@ -29,6 +34,9 @@
         }
     )
 )
+_MARKDOWN_IMPORT_DIALECTS: Mapping[
+    MarkdownFlavorName, MarkdownImportDialect
+] = MappingProxyType({_MYSTMD_MARKDOWN_IMPORT.name: _MYSTMD_MARKDOWN_IMPORT})
 # Filename inference handles target-specific markdown extensions.
 _MARKDOWN_FLAVORS_BY_EXTENSION: Mapping[str, MarkdownFlavor] = (
     MappingProxyType({".myst.md": _MYSTMD_MARKDOWN, ".qmd": _QMD_MARKDOWN})
@@ -81,6 +89,16 @@ def normalize_markdown_flavor(
         raise ValueError(f"Unsupported markdown flavor: {flavor!r}") from error
 
 
+def _markdown_import_dialects(
+    text: str, filepath: str | None
+) -> tuple[MarkdownImportDialect, ...]:
+    return tuple(
+        dialect
+        for dialect in _MARKDOWN_IMPORT_DIALECTS.values()
+        if dialect.matches(text, filepath)
+    )
+
+
 def _markdown_output_extension(
     flavor: MarkdownFlavor | MarkdownFlavorName,
 ) -> str:

marimo/_convert/markdown/flavor/base.py

@@ -2,8 +2,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar, Literal
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, ClassVar, Literal, Protocol
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -42,6 +42,29 @@ class MarkdownExportDocument:
     blocks: list[MarkdownExportBlock]
 
 
+@dataclass
+class MarkdownImportContext:
+    """Mutable state shared by markdown import dialects."""
+
+    metadata: dict[str, str] = field(default_factory=dict)
+
+
+class MarkdownImportDialect(Protocol):
+    """Source markdown syntax adapter for the canonical importer."""
+
+    name: MarkdownFlavorName
+
+    def matches(self, text: str, filepath: str | None) -> bool:
+        """Return whether this dialect should preprocess the markdown."""
+        ...
+
+    def preprocess(
+        self, lines: list[str], context: MarkdownImportContext
+    ) -> list[str]:
+        """Normalize source markdown before the canonical importer runs."""
+        ...
+
+
 class MarkdownFlavor(ABC):
     """Markdown-family output flavor.
 

marimo/_convert/markdown/flavor/mystmd.py

@@ -15,27 +15,96 @@
 import re
 from typing import TYPE_CHECKING
 
+from marimo import _loggers
 from marimo._convert.markdown.flavor.base import (
     CodeCellBlock,
     MarkdownCellBlock,
     MarkdownExportDocument,
     MarkdownFlavor,
+    MarkdownFlavorName,
+    MarkdownImportContext,
+    _escape_attribute,
 )
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
+LOGGER = _loggers.marimo_logger()
+
 # Metadata emitted through the `{marimo-config}` directive.
 _CONFIG_KEYS = {"header", "pyproject"}
 # marimo-specific metadata filtered before writing MyST frontmatter.
 _MARIMO_METADATA_KEYS = {"width"}
+# MyST marimo executable directive headers.
+_MARIMO_DIRECTIVE_HEADER_RE = re.compile(
+    r"^(?P<fence>`{3,})\{marimo\}(?:\s+(?P<language>\w+))?\s*$"
+)
+# MyST marimo page-level configuration directive headers.
+_MARIMO_CONFIG_HEADER_RE = re.compile(
+    r"^(?P<fence>`{3,})\{marimo-config\}\s*$"
+)
+# MyST directive options, e.g. `:hide-code: true`.
+_DIRECTIVE_OPTION_RE = re.compile(r"^:([A-Za-z0-9_-]+):(?:\s+(.*))?$")
 # PEP 723 script metadata blocks embedded in exported notebook headers.
 _SCRIPT_METADATA_RE = re.compile(
     r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s"
     r"(?P<content>(^#(| .*)$\s)+)^# ///$"
 )
 
 
+class _MystmdMarkdownImportDialect:
+    name: MarkdownFlavorName = "mystmd"
+
+    def matches(self, text: str, filepath: str | None) -> bool:
+        del filepath
+        return any(_is_marimo_header(line) for line in text.splitlines())
+
+    def preprocess(
+        self, lines: list[str], context: MarkdownImportContext
+    ) -> list[str]:
+        normalized: list[str] = []
+        index = 0
+
+        while index < len(lines):
+            config_match = _MARIMO_CONFIG_HEADER_RE.match(lines[index])
+            if config_match is not None:
+                closing_index = _find_closing_fence(
+                    lines, index + 1, config_match.group("fence")
+                )
+                if closing_index is None:
+                    normalized.extend(lines[index:])
+                    break
+
+                context.metadata.update(
+                    _extract_config_metadata(lines[index + 1 : closing_index])
+                )
+                index = closing_index + 1
+                continue
+
+            match = _MARIMO_DIRECTIVE_HEADER_RE.match(lines[index])
+            if match is None:
+                normalized.append(lines[index])
+                index += 1
+                continue
+
+            closing_index = _find_closing_fence(
+                lines, index + 1, match.group("fence")
+            )
+            if closing_index is None:
+                normalized.extend(lines[index:])
+                break
+
+            options, body_lines = _extract_directive_options(
+                lines[index + 1 : closing_index]
+            )
+            normalized.append(_canonical_code_fence_head(match, options))
+            normalized.extend(body_lines)
+            normalized.append(lines[closing_index])
+            index = closing_index + 1
+
+        return normalized
+
+
 class MystmdMarkdownFlavor(MarkdownFlavor):
     """Render marimo notebooks as mystmd markdown.
 
@@ -164,3 +233,82 @@ def _uncomment_script_metadata(content: str) -> str:
 
 def _mystmd_option_name(key: str) -> str:
     return key.replace("_", "-")
+
+
+def _is_marimo_header(line: str) -> bool:
+    return bool(
+        _MARIMO_DIRECTIVE_HEADER_RE.match(line)
+        or _MARIMO_CONFIG_HEADER_RE.match(line)
+    )
+
+
+def _is_closing_fence(line: str, opening_fence: str) -> bool:
+    stripped = line.strip()
+    return len(stripped) >= len(opening_fence) and set(stripped) == {"`"}
+
+
+def _find_closing_fence(
+    lines: list[str], start: int, opening_fence: str
+) -> int | None:
+    for index in range(start, len(lines)):
+        if _is_closing_fence(lines[index], opening_fence):
+            return index
+    return None
+
+
+def _extract_directive_options(
+    lines: list[str],
+) -> tuple[dict[str, str], list[str]]:
+    options: dict[str, str] = {}
+    body_start = 0
+
+    for index, line in enumerate(lines):
+        match = _DIRECTIVE_OPTION_RE.match(line)
+        if match is None:
+            break
+        options[match.group(1).replace("-", "_")] = match.group(2) or "true"
+        body_start = index + 1
+
+    if body_start and body_start < len(lines) and lines[body_start] == "":
+        body_start += 1
+
+    return options, lines[body_start:]
+
+
+def _canonical_code_fence_head(
+    match: re.Match[str], options: dict[str, str]
+) -> str:
+    attribute_str = "".join(
+        f' {key}="{_escape_attribute(value)}"'
+        for key, value in options.items()
+    )
+    return "{fence}{language} {{.marimo{attributes}}}".format(
+        fence=match.group("fence"),
+        language=match.group("language") or "python",
+        attributes=attribute_str,
+    )
+
+
+def _extract_config_metadata(lines: list[str]) -> dict[str, str]:
+    from marimo._utils import yaml
+
+    if lines and lines[0] == "---":
+        for index, line in enumerate(lines[1:], start=1):
+            if line == "---":
+                lines = lines[1:index]
+                break
+
+    try:
+        metadata = yaml.load("\n".join(lines))
+    except yaml.YAMLError:
+        LOGGER.warning("Error parsing marimo-config YAML. Ignoring config.")
+        return {}
+
+    if not isinstance(metadata, dict):
+        return {}
+
+    return {
+        key: value
+        for key, value in metadata.items()
+        if key in _CONFIG_KEYS and isinstance(value, str)
+    }