Add support for Marimo notebooks

docs/contents.md

@@ -33,6 +33,8 @@ pip install mols2grid
 
 mols2grid is mainly meant to be used in notebooks (Jupyter notebooks, Jupyter Lab, and Google Colab) but it can also be used as a standalone HTML page opened with your favorite web browser, or embedded in a Streamlit app.
 
+Since v2.2.0, mols2grid is also compatible with [Marimo](https://marimo.io/). An example is available in [`scripts/marimo_example.py`](https://github.com/cbouy/mols2grid/blob/master/scripts/marimo_example.py).
+
 Since Streamlit doesn't seem to support ipywidgets yet, some features aren't functional: retrieving the selection from Python (you can still export it from the GUI) and using Python callbacks.
 
 <img alt="knime logo" align="left" style="padding:6px" src="https://www.knime.com/sites/default/files/favicons/favicon-32x32.png"/>
@@ -84,6 +86,7 @@ Feel free to open a pull request if you'd like your snippets to be added to this
 * [@fredrikw](https://github.com/fredrikw) (contributor)
 * [@JustinChavez](https://github.com/JustinChavez) (contributor)
 * [@hadim](https://github.com/hadim) (conda feedstock maintainer)
+* [@N283T](https://github.com/N283T) (contributor)
 
 # 🎓 Citing
 ---

mols2grid/molgrid.py

@@ -1,3 +1,4 @@
+import ast
 import json
 import warnings
 from base64 import b64encode
@@ -14,6 +15,7 @@
 from mols2grid.utils import (
     callback_handler,
     env,
+    is_running_within_marimo,
     is_running_within_streamlit,
     mol_to_record,
     mol_to_smiles,
@@ -205,7 +207,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
@@ -679,10 +682,10 @@ def to_interactive(  # noqa: PLR0912
         # Generate cell HTML.
         item = (
             '<div class="m2g-cell" data-mols2grid-id="0" tabindex="0">'
-            '<div class="m2g-cb-wrap">{checkbox_html}<div class="m2g-cb"></div>'
-            "{id_display_html}</div>"
-            '<div class="m2g-cell-actions">{info_btn_html}{callback_btn_html}</div>'
-            "{content}"
+            '<div class="m2g-cb-wrap">{checkbox_html}<div class="m2g-cb"></div>'  # noqa: RUF027
+            "{id_display_html}</div>"  # noqa: RUF027
+            '<div class="m2g-cell-actions">{info_btn_html}{callback_btn_html}</div>'  # noqa: RUF027
+            "{content}"  # noqa: RUF027
             "{tooltip_html}"
             "</div>"
         )
@@ -800,6 +803,36 @@ def get_selection(self):
             columns=self._extra_columns
         )
 
+    def get_marimo_selection(self):
+        """Returns a marimo state object containing the list of selected indices.
+        Only available when running in marimo.
+
+        Returns
+        -------
+        getter
+            A getter function for the selection state.
+            Calling it with no arguments returns the current list of selected IDs.
+        """
+        if not is_running_within_marimo():
+            raise RuntimeError("This method is only available in a marimo notebook.")
+
+        import marimo as mo
+
+        get_state, set_state = mo.state([])
+
+        def _on_change(change):
+            try:
+                sel = ast.literal_eval(change["new"])
+                set_state(list(sel.keys()))
+            except (ValueError, SyntaxError):
+                pass
+
+        if not getattr(self.widget, "_marimo_hooked", False):
+            self.widget.observe(_on_change, names=["selection"])
+            self.widget._marimo_hooked = True
+
+        return get_state
+
     def filter(self, mask):
         """Filters the grid using a mask (boolean array).
 
@@ -1073,8 +1106,13 @@ def display(
         -------
         view : IPython.core.display.HTML
         """
+        requires_marimo = is_running_within_marimo()
+        if requires_marimo:
+            use_iframe = True
+
         use_iframe = is_jupyter or use_iframe
         doc = self.render(**kwargs, use_iframe=use_iframe)
+
         if use_iframe:
             # Render HTML in iframe.
             iframe = env.get_template("html/iframe.html").render(
@@ -1084,6 +1122,10 @@ def display(
                 sandbox=iframe_sandbox,
                 doc=escape(doc),
             )
+            if requires_marimo:
+                import marimo as mo
+
+                return mo.vstack([self.widget, mo.Html(iframe)])
             return HTML(iframe)
         # Render HTML regularly.
         return HTML(doc)

mols2grid/select.py

@@ -1,6 +1,8 @@
 import warnings
 from ast import literal_eval
 
+from mols2grid.utils import is_running_within_marimo
+
 
 class SelectionRegister:
     """Register for grid selections
@@ -21,7 +23,7 @@ def _update_current_grid(self, name):
 
     def _init_grid(self, name):
         overwrite = self.SELECTIONS.get(name, False)
-        if overwrite:
+        if overwrite and not is_running_within_marimo():
             warnings.warn(
                 f"Overwriting non-empty {name!r} grid selection: {overwrite!s}",
                 stacklevel=2,

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,8 @@ dynamic = ["version"]
 [project.license]
 file = "LICENSE"
 
+
+
 [dependency-groups]
 build = ["build", "watchfiles", "jupyterlab"]
 tests = [
@@ -46,6 +48,7 @@ tests = [
     "cairosvg~=2.5",
     "imagehash~=4.3",
     "selenium~=4.11",
+    "marimo>=0.18.4",
 ]
 docs = [
     "py3dmol",
@@ -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"

scripts/marimo_example.py

@@ -0,0 +1,103 @@
+import marimo
+
+__generated_with = "0.18.4"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def import_libraries():
+    import marimo as mo
+
+    import mols2grid
+    from mols2grid.datafiles import SOLUBILITY_SDF
+
+    return SOLUBILITY_SDF, mo, mols2grid
+
+
+@app.cell
+def prepare(mo):
+    from rdkit.Chem.Draw import rdMolDraw2D
+
+    def mol_to_svg(mol, width=130, height=90, opts=None):
+        if mol is None:
+            return mo.Html("")
+
+        drawer = rdMolDraw2D.MolDraw2DSVG(width, height)
+
+        if opts is not None:
+            drawer.SetDrawOptions(opts)
+
+        drawer.DrawMolecule(mol)
+        drawer.FinishDrawing()
+        return mo.Html(drawer.GetDrawingText())
+
+    solubility_range = mo.ui.range_slider(
+        -10,
+        2,
+        0.5,
+        debounce=True,
+        show_value=True,
+        full_width=True,
+        label="Solubility",
+    )
+    return mol_to_svg, solubility_range
+
+
+@app.cell
+def create_grid(SOLUBILITY_SDF, mols2grid):
+    # NOTE:
+    # This cell is intentionally kept independent from the sliders.
+    # In marimo, cells are re-executed whenever any of their dependencies change.
+    # Keeping grid creation here prevents MolGrid.from_sdf(...) from being
+    # re-run on every slider update, which would reset the widget state.
+    grid = mols2grid.MolGrid.from_sdf(SOLUBILITY_SDF, size=(120, 100))
+    get_selection_ids = grid.get_marimo_selection()
+    view = grid.display(n_items_per_page=12, selection=True)
+    return get_selection_ids, grid, view
+
+
+@app.cell
+def filter_and_display(grid, mo, solubility_range, view):
+    mask = grid.dataframe["SOL"].between(*solubility_range.value)
+    results = grid.dataframe.loc[mask]
+    # Same as:
+    # grid.dataframe["SOL"] >= solubility_range.value[0]) & \
+    #   (grid.dataframe["SOL"] <= solubility_range.value[1])
+
+    grid.filter_by_index(results.index)
+
+    mo.vstack([solubility_range, view])
+    return (results,)
+
+
+@app.cell(hide_code=True)
+def display_selection(get_selection_ids, mo, mol_to_svg, results):
+    # This cell displays the selected molecules in a Marimo table.
+    # The `mol_to_svg` function is used to render the molecule images
+    # directly in the table.
+    # We filter the dataframe based on the selection state (`get_selection_ids`)
+    # from the grid above.
+
+    selected = results[results["mols2grid-id"].isin(get_selection_ids())]
+
+    # # If you want to use custom drawing options:
+    # opts = rdMolDraw2D.MolDrawOptions()
+    # opts.explicitMethyl = True
+
+    table = mo.ui.table(
+        selected.reset_index(drop=True).drop(columns="img"),
+        format_mapping={
+            "mol": mol_to_svg
+            # "mol": lambda mol: mol_to_svg(mol, opts=opts)
+            # If you want to use custom drawing options
+        },
+        freeze_columns_left=["mols2grid-id", "mol"],
+        freeze_columns_right=["SOL"],
+        label="Try selecting molecules from the grid above!!",
+    )
+
+    table  # noqa: B018
+
+
+if __name__ == "__main__":
+    app.run()