fix: execution order for marimo

Author: cbouy

mols2grid/molgrid.py

@@ -804,6 +804,10 @@ def get_marimo_selection(self):
         """
         if not is_running_within_marimo():
             raise RuntimeError("This method is only available in a marimo notebook.")
+        if not hasattr(self, "widget"):
+            raise RuntimeError(
+                "Please run the `display` method first to render the underlying widget"
+            )
 
         import marimo as mo
 
@@ -1095,12 +1099,16 @@ def display(
         -------
         view : IPython.core.display.HTML
         """
-        use_iframe = is_jupyter if use_iframe is None else use_iframe
         obj = self.render(**kwargs, use_iframe=use_iframe)
-
+        # if widget
         if not isinstance(obj, str):
+            if is_running_within_marimo():
+                import marimo as mo
+
+                return mo.ui.anywidget(obj)
             return obj
 
+        use_iframe = is_jupyter if use_iframe is None else use_iframe
         if not use_iframe:
             return HTML(obj)
 

scripts/marimo_example.py

@@ -51,8 +51,8 @@ def create_grid(SOLUBILITY_SDF, mols2grid):
     # 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)
+    get_selection_ids = grid.get_marimo_selection()
     return get_selection_ids, grid, view
 
 
@@ -95,7 +95,6 @@ def display_selection(get_selection_ids, mo, mol_to_svg, results):
         freeze_columns_right=["SOL"],
         label="Try selecting molecules from the grid above!!",
     )
-
     table  # noqa: B018