feat: add inline equation source editing helper

examples/notebooks/spec_black_scholes_with_results.py

@@ -1,12 +1,11 @@
 import marimo
 
-__generated_with = "0.22.4"
+__generated_with = "0.23.1"
 app = marimo.App(width="medium")
 
 
 @app.cell
 def notebook_context():
-
     import marimo as mo
 
     from pdealchemy.exceptions import PDEAlchemyError
@@ -19,7 +18,7 @@ def notebook_context():
         render_notebook_report_controls,
         selection_from_controls,
     )
-    from pdealchemy.notebook_utils import math_eq, spec_md
+    from pdealchemy.notebook_utils import math_eq, math_eq_editor, spec_md
 
     pdealchemy_error = PDEAlchemyError
     return (
@@ -29,6 +28,7 @@ def notebook_context():
         compose_report_view,
         create_notebook_report_controls,
         math_eq,
+        math_eq_editor,
         mo,
         pdealchemy_error,
         render_notebook_report_controls,
@@ -89,8 +89,8 @@ def sde(math_eq):
 
 
 @app.cell
-def pde(math_eq):
-    math_eq("library/pde/black_scholes.md", name="Main PDE operator")
+def pde(math_eq_editor):
+    math_eq_editor("library/pde/black_scholes.md", name="Main PDE operator")
     return
 
 
@@ -184,12 +184,7 @@ def _(
 
 
 @app.cell
-def _(
-    build_report_table_views,
-    mo,
-    report,
-    selection,
-):
+def _(build_report_table_views, mo, report, selection):
     table_views = build_report_table_views(
         mo,
         report=report,
@@ -199,11 +194,7 @@ def _(
 
 
 @app.cell
-def _(
-    build_report_chart_views,
-    report,
-    selection,
-):
+def _(build_report_chart_views, report, selection):
     chart_views = build_report_chart_views(
         report=report,
         selection=selection,
@@ -212,7 +203,7 @@ def _(
 
 
 @app.cell
-def _(chart_views: list[object], compose_report_view, error_view, mo, table_views: list[object]):
+def _(chart_views, compose_report_view, error_view, mo, table_views):
     view = compose_report_view(
         mo,
         error_view=error_view,

src/pdealchemy/notebook_spec.py

@@ -20,7 +20,7 @@
     "boundary_upper": ("boundary", "upper"),
     "discretisation": ("numerics",),
 }
-_SUPPORTED_EQUATION_HELPERS = {"math_eq", "eq_from_file"}
+_SUPPORTED_EQUATION_HELPERS = {"math_eq", "math_eq_editor", "eq_from_file"}
 _SUPPORTED_MARKDOWN_HELPERS = {"mo.md"}
 _SUPPORTED_FILE_MARKDOWN_HELPERS = {"spec_md"}
 

src/pdealchemy/notebook_utils.py

@@ -38,15 +38,27 @@ def _render_markdown_block(*, markdown: str, name: str | None, mo: ModuleType) -
     return mo.md(f"### {name.strip()}\n\n{markdown}")
 
 
+def _extract_first_latex_block(text: str) -> str:
+    """Extract the first LaTeX display block from markdown-like content."""
+    match = re.search(r"\\\[(.+?)\\\]", text, re.DOTALL)
+    if match is None:
+        return text.strip()
+    return match.group(1).strip()
+
+
+def _read_equation_source(source_path: Path) -> str:
+    """Read equation source content from disk."""
+    return source_path.read_text(encoding="utf-8")
+
+
 def math_eq(content: str, *, name: str | None = None) -> object:
     """Render LaTeX content or file-backed LaTeX with an optional equation heading."""
     mo = _load_marimo_module()
     candidate_path = Path(content).expanduser()
     if candidate_path.is_file():
         try:
-            text = candidate_path.read_text(encoding="utf-8").strip()
-            match = re.search(r"\\\[(.+?)\\\]", text, re.DOTALL)
-            latex = match.group(1).strip() if match else text
+            source_text = _read_equation_source(candidate_path)
+            latex = _extract_first_latex_block(source_text)
             return _render_equation_block(latex=latex, name=name, mo=mo)
         except OSError as exc:
             return mo.md(f"**Error loading equation file** `{content}`: {exc}")
@@ -64,3 +76,62 @@ def spec_md(content: str, *, name: str | None = None) -> object:
         except OSError as exc:
             return mo.md(f"**Error loading markdown file** `{content}`: {exc}")
     return _render_markdown_block(markdown=content, name=name, mo=mo)
+
+
+def math_eq_editor(
+    content: str,
+    *,
+    name: str | None = None,
+    save_label: str = "Save equation file",
+) -> object:
+    """Render an inline markdown equation editor with save and preview controls."""
+    mo = _load_marimo_module()
+    source_path = Path(content).expanduser()
+    if not source_path.is_file():
+        return mo.md(
+            f"**Error loading equation file** `{content}`: file not found. "
+            "Provide a valid markdown equation path."
+        )
+    try:
+        source_text = _read_equation_source(source_path)
+    except OSError as exc:
+        return mo.md(f"**Error loading equation file** `{content}`: {exc}")
+
+    source_editor = mo.ui.code_editor(
+        value=source_text,
+        language="markdown",
+        label=f"Equation source: {content}",
+    )
+
+    def _save_source(_value: object) -> str:
+        try:
+            source_path.write_text(str(source_editor.value), encoding="utf-8")
+        except OSError as exc:
+            return f"Save failed: {exc}"
+        return f"Saved `{content}`."
+
+    save_button = mo.ui.button(
+        on_click=_save_source,
+        label=save_label,
+        kind="success",
+    )
+    save_status = save_button.value
+    status_message = (
+        f"_Save status_: {save_status}"
+        if save_status is not None
+        else "_Edit the source and click Save equation file to persist changes._"
+    )
+    preview = _render_equation_block(
+        latex=_extract_first_latex_block(str(source_editor.value)),
+        name=name,
+        mo=mo,
+    )
+    return mo.vstack(
+        [
+            mo.md(f"### Inline equation source editor\n\n`{content}`"),
+            source_editor,
+            save_button,
+            mo.md(status_message),
+            preview,
+        ]
+    )

tests/notebook/test_notebook_spec.py

@@ -100,6 +100,31 @@ def test_notebook_to_toml_supports_app_cell_call_decorator(tmp_path: Path) -> No
     assert 'equation_file = "library/pde/black_scholes.md"' in rendered
 
 
+def test_notebook_to_toml_supports_math_eq_editor_helper(tmp_path: Path) -> None:
+    notebook_path = tmp_path / "spec_math_eq_editor.py"
+    _write_notebook(
+        notebook_path,
+        [
+            "import marimo as mo",
+            "from pdealchemy.notebook_utils import math_eq_editor",
+            "",
+            "app = mo.App()",
+            "",
+            'mo.md("# Editor Helper")',
+            "",
+            "@app.cell",
+            "def pde():",
+            '    """Main PDE operator."""',
+            '    math_eq_editor("library/pde/black_scholes.md", name="Main PDE operator")',
+        ],
+    )
+
+    rendered = notebook_to_toml_content(notebook_path)
+
+    assert "[mathematics.operator]" in rendered
+    assert 'equation_file = "library/pde/black_scholes.md"' in rendered
+
+
 def test_notebook_to_toml_ignores_unknown_cells(tmp_path: Path) -> None:
     notebook_path = tmp_path / "spec_mixed_cells.py"
     _write_notebook(

tests/notebook/test_notebook_utils.py

@@ -17,6 +17,54 @@ def _fake_marimo_module() -> SimpleNamespace:
     return SimpleNamespace(md=lambda text: text)
 
 
+class _FakeCodeEditor:
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+
+class _FakeButton:
+    def __init__(self, on_click: Any) -> None:
+        self._on_click = on_click
+        self.value: Any = None
+
+    def click(self) -> None:
+        self.value = self._on_click(None)
+
+
+class _FakeUi:
+    def code_editor(
+        self,
+        value: str,
+        *,
+        language: str,
+        label: str,
+    ) -> _FakeCodeEditor:
+        _ = (language, label)
+        return _FakeCodeEditor(value)
+
+    def button(
+        self,
+        *,
+        on_click: Any,
+        label: str,
+        kind: str,
+    ) -> _FakeButton:
+        _ = (label, kind)
+        return _FakeButton(on_click)
+
+
+class _FakeEditorMarimo:
+    ui = _FakeUi()
+
+    @staticmethod
+    def md(text: str) -> tuple[str, str]:
+        return ("md", text)
+
+    @staticmethod
+    def vstack(blocks: list[object]) -> tuple[str, list[object]]:
+        return ("vstack", blocks)
+
+
 def test_math_eq_renders_raw_latex(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(notebook_utils, "_load_marimo_module", _fake_marimo_module)
 
@@ -118,6 +166,64 @@ def _raise_read_error(*_args: Any, **_kwargs: Any) -> str:
     assert "Cannot read file" in rendered
 
 
+def test_math_eq_editor_renders_inline_editor_with_preview(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    equation_file = tmp_path / "equation.md"
+    equation_file.write_text(
+        "\n".join(
+            [
+                "PDE notes",
+                "\\[",
+                "\\frac{\\partial V}{\\partial t} + rSV_S = 0",
+                "\\]",
+            ]
+        ),
+        encoding="utf-8",
+    )
+    monkeypatch.setattr(notebook_utils, "_load_marimo_module", lambda: _FakeEditorMarimo)
+
+    rendered = notebook_utils.math_eq_editor(str(equation_file), name="Main PDE operator")
+
+    assert rendered[0] == "vstack"
+    editor = next(block for block in rendered[1] if isinstance(block, _FakeCodeEditor))
+    preview = rendered[1][-1]
+    assert "PDE notes" in editor.value
+    assert preview[0] == "md"
+    assert "### Main PDE operator" in preview[1]
+    assert "\\frac{\\partial V}{\\partial t} + rSV_S = 0" in preview[1]
+
+
+def test_math_eq_editor_persists_edited_source(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    equation_file = tmp_path / "equation.md"
+    equation_file.write_text("\\[\nS\n\\]", encoding="utf-8")
+    monkeypatch.setattr(notebook_utils, "_load_marimo_module", lambda: _FakeEditorMarimo)
+
+    rendered = notebook_utils.math_eq_editor(str(equation_file))
+    editor = next(block for block in rendered[1] if isinstance(block, _FakeCodeEditor))
+    button = next(block for block in rendered[1] if isinstance(block, _FakeButton))
+    editor.value = "\\[\nS-K\n\\]"
+
+    button.click()
+
+    assert equation_file.read_text(encoding="utf-8") == "\\[\nS-K\n\\]"
+    assert button.value == f"Saved `{equation_file}`."
+
+
+def test_math_eq_editor_rejects_missing_file(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(notebook_utils, "_load_marimo_module", lambda: _FakeEditorMarimo)
+
+    rendered = notebook_utils.math_eq_editor("library/pde/missing.md")
+
+    assert rendered[0] == "md"
+    assert "Error loading equation file" in rendered[1]
+    assert "file not found" in rendered[1]
+
+
 def test_load_marimo_module_raises_config_error_when_missing(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: