Add support for Marimo notebooks

mols2grid/molgrid.py

@@ -14,6 +14,7 @@
 from mols2grid.utils import (
     callback_handler,
     env,
+    is_running_within_marimo,
     is_running_within_streamlit,
     mol_to_record,
     mol_to_smiles,
@@ -205,7 +206,8 @@ def __init__(  # noqa: PLR0912
         widget.observe(selection_handler, names=["selection"])
 
         # Register widget JS-side.
-        display(widget)
+        if not is_running_within_marimo():
+            display(widget)
         self.widget = widget
 
     @classmethod
@@ -1073,8 +1075,27 @@ def display(
         -------
         view : IPython.core.display.HTML
         """
+        if is_running_within_marimo():
+            use_iframe = True
+
         use_iframe = is_jupyter or use_iframe
         doc = self.render(**kwargs, use_iframe=use_iframe)
+
+        if is_running_within_marimo():
+            import marimo as mo
+
+            if use_iframe:
+                # Render HTML in iframe.
+                iframe = env.get_template("html/iframe.html").render(
+                    width=iframe_width,
+                    height=iframe_height,
+                    allow=iframe_allow,
+                    sandbox=iframe_sandbox,
+                    doc=escape(doc),
+                )
+                return mo.vstack([self.widget, mo.Html(iframe)])
+            return mo.vstack([self.widget, mo.Html(doc)])
+
         if use_iframe:
             # Render HTML in iframe.
             iframe = env.get_template("html/iframe.html").render(

mols2grid/utils.py

@@ -129,3 +129,17 @@ def is_running_within_streamlit():
         return False
     else:
         return ctx is not None
+
+
+def is_running_within_marimo():
+    """
+    Function to check whether python code is run within marimo
+
+    Returns
+    -------
+    use_marimo : boolean
+        True if code is run within marimo, else False
+    """
+    import sys
+
+    return "marimo" in sys.modules

pyproject.toml

@@ -37,6 +37,9 @@ dynamic = ["version"]
 [project.license]
 file = "LICENSE"
 
+[project.optional-dependencies]
+marimo = ["marimo>=0.18.4"]
+
 [dependency-groups]
 build = ["build", "watchfiles", "jupyterlab"]
 tests = [
@@ -235,4 +238,4 @@ sequence = [{ ref = "style-check" }, { ref = "tests" }, { ref = "docs" }]
 
 [tool.poe.tasks.build]
 help = "Builds the package"
-cmd = "python -m build"
+cmd = "python -m build"

reproduce_marimo.py

@@ -0,0 +1,54 @@
+import marimo
+
+__generated_with = "0.18.4"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def _():
+    import marimo as mo
+    from rdkit import Chem
+    import mols2grid
+    import pandas as pd
+    from io import StringIO
+    return StringIO, mols2grid, pd
+
+
+@app.cell
+def _():
+    smiles = """SMILES	NAME
+    CCO	Ethanol
+    CC(=O)O	Acetic_acid
+    CCC	Propane
+    C1=CC=CC=C1	Benzene
+    COC	Dimethyl_ether
+    CCN	Ethylamine
+    C(CO)O	Ethylene_glycol
+    CCOCC	Diethyl_ether
+    OC=O	Formic_acid
+    CCOC(=O)C	Ethyl_acetate
+    """
+    return (smiles,)
+
+
+@app.cell
+def _(StringIO, pd, smiles):
+    df = pd.read_csv(StringIO(smiles), delimiter="\t")
+    df
+    return (df,)
+
+
+@app.cell
+def _(df, mols2grid):
+    mols2grid.display(df)
+    return
+
+
+@app.cell
+def _(mols2grid):
+    mols2grid.get_selection()
+    return
+
+
+if __name__ == "__main__":
+    app.run()

tests/test_marimo_integration.py

@@ -0,0 +1,58 @@
+
+import sys
+import pytest
+from unittest.mock import patch, MagicMock
+import pandas as pd
+from mols2grid import MolGrid
+from mols2grid.utils import is_running_within_marimo
+
+@pytest.fixture
+def mock_marimo_module():
+    with patch.dict(sys.modules, {"marimo": MagicMock()}):
+        yield
+
+def test_is_running_within_marimo_true(mock_marimo_module):
+    assert is_running_within_marimo() is True
+
+def test_is_running_within_marimo_false():
+    # Ensure marimo is not in sys.modules for this test
+    with patch.dict(sys.modules):
+        if "marimo" in sys.modules:
+            del sys.modules["marimo"]
+        assert is_running_within_marimo() is False
+
+def test_init_in_marimo_does_not_display(mock_marimo_module):
+    df = pd.DataFrame({"SMILES": ["C"]})
+
+    # Mock IPython.display.display which is imported as display in molgrid.py
+    # We need to patch it where it is used, i.e., in mols2grid.molgrid
+    with patch("mols2grid.molgrid.display") as mock_display:
+        mg = MolGrid(df, smiles_col="SMILES")
+        mock_display.assert_not_called()
+
+def test_display_in_marimo(mock_marimo_module):
+    df = pd.DataFrame({"SMILES": ["C"]})
+    mg = MolGrid(df, smiles_col="SMILES")
+
+    # Mock marimo.Html and marimo.vstack
+    with patch("marimo.Html") as mock_html, \
+         patch("marimo.vstack") as mock_vstack:
+
+        result = mg.display()
+
+        # Verify that an iframe is being rendered inside Html
+        mock_html.assert_called_once()
+        args, _ = mock_html.call_args
+        html_content = args[0]
+        assert "<iframe" in html_content
+        assert 'class="mols2grid-iframe"' in html_content
+
+        # Verify vstack was called with [widget, html]
+        mock_vstack.assert_called_once()
+        vstack_args = mock_vstack.call_args[0][0]
+        assert len(vstack_args) == 2
+        assert vstack_args[0] == mg.widget
+        assert vstack_args[1] == mock_html.return_value
+
+        # Ensure the result is the return value of marimo.vstack
+        assert result == mock_vstack.return_value