refactor: register

Author: cbouy

mols2grid/__init__.py

@@ -4,7 +4,7 @@
 from mols2grid.callbacks import make_popup_callback
 from mols2grid.dispatch import display, save
 from mols2grid.molgrid import MolGrid
-from mols2grid.select import get_selection, list_grids
+from mols2grid.select import get_selection, link_marimo_state, list_grids
 from mols2grid.utils import is_running_within_streamlit, sdf_to_dataframe
 
 if is_running_within_streamlit():

mols2grid/molgrid.py

@@ -1,4 +1,3 @@
-import ast
 import json
 import os
 import warnings
@@ -640,9 +639,9 @@ def to_interactive(  # noqa: PLR0912
             try:
                 cached_selection = register.get_selection(self._grid_id)
             except KeyError:
-                register._init_grid(self._grid_id)
+                register.add_grid(self._grid_id)
             else:
-                register._update_current_grid(self._grid_id)
+                register.current_selection = self._grid_id
             if self.cache_selection:
                 df["cached_checkbox"] = False
                 df.loc[
@@ -785,40 +784,6 @@ 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.")
-        if not hasattr(self, "widget"):
-            raise RuntimeError(
-                "Please run the `display` method first to render the underlying widget"
-            )
-
-        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).
 

mols2grid/select.py

@@ -17,47 +17,110 @@ class SelectionRegister:
 
     def __init__(self):
         self.SELECTIONS = {}
+        self.CALLBACKS = []
+        self._current_selection = None
 
-    def _update_current_grid(self, name):
-        self.current_selection = name
+    @property
+    def current_selection(self):
+        """The name of the last updated grid (created or interacted with),
+        or ``None`` if no grid has been displayed yet."""
+        return self._current_selection
+
+    @current_selection.setter
+    def current_selection(self, name):
+        if name is not None and name not in self.SELECTIONS:
+            raise ValueError(
+                f"The selection for {name} must be initialized "
+                "before setting it as the current grid"
+            )
+        self._current_selection = name
 
-    def _init_grid(self, name):
+    @current_selection.deleter
+    def current_selection(self):
+        self._current_selection = None
+
+    def add_grid(self, name):
+        """Adds a grid to track selections for."""
         overwrite = self.SELECTIONS.get(name, False)
         if overwrite and not is_running_within_marimo():
             warnings.warn(
                 f"Overwriting non-empty {name!r} grid selection: {overwrite!s}",
                 stacklevel=2,
             )
         self.SELECTIONS[name] = {}
-        self._update_current_grid(name)
+        self.current_selection = name
+        self.add_callback(self._store_selection)
 
     def selection_updated(self, name, event):
-        self.SELECTIONS[name] = literal_eval(event.new)
-        self._update_current_grid(name)
+        """Callback function linked to the widget."""
+        self.current_selection = name
+        selection = literal_eval(event.new)
+        for callback in self.CALLBACKS:
+            callback(name, selection)
+
+    def _store_selection(self, name, selection):
+        """Makes the selection available to the register."""
+        self.SELECTIONS[name] = selection
+
+    def add_callback(self, callback):
+        """Add a callback function to be called when the selection is updated.
+
+        Parameters
+        ----------
+        callback : callable
+            The function to execute when a the selection is updated.
+        """
+        self.CALLBACKS.append(callback)
 
     def get_selection(self, name=None):
         """Returns the selection for a specific MolGrid instance
 
         Parameters
         ----------
         name : str or None
-            Name of the grid to fetch the selection from. If `None`, the most
+            Name of the grid to fetch the selection from. If ``None``, the most
             recently updated grid is returned
         """
         name = self.current_selection if name is None else name
         return self.SELECTIONS[name]
 
+    def link_marimo_state(self):
+        """Link the register to marimo by initializing a ``state`` dict.
+        When the selection on a grid is updated, the state setter function from marimo
+        is called to update the state's value.
+
+        Returns
+        -------
+        get_state, set_state:
+            The state getter and setter returned by ``marimo.state({})``.
+            ``get_state`` returns a dictionary of all selections where the keys are
+            grid names and values are the indices selected for that grid.
+        """
+        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 marimo_callback(name, selection):
+            set_state(lambda value: {**value, name: list(selection)})
+
+        self.add_callback(marimo_callback)
+        return get_state, set_state
+
     def list_grids(self):
         """Returns a list of grid names"""
-        return list(self.SELECTIONS.keys())
+        return list(self.SELECTIONS)
 
-    def _clear(self):
+    def clear(self):
         """Clears all selections"""
-        if hasattr(self, "current_selection"):
-            del self.current_selection
         self.SELECTIONS.clear()
+        self.CALLBACKS.clear()
+        del self.current_selection
 
 
 register = SelectionRegister()
 get_selection = register.get_selection
 list_grids = register.list_grids
+link_marimo_state = register.link_marimo_state

tests/test_interface.py

@@ -176,7 +176,7 @@ def test_selection_click(driver: CustomDriver, html_doc):
     driver.wait_for_img_load()
     sel = driver.click_checkbox()
     assert sel == {0: "CCC(C)CC"}
-    register._clear()
+    register.clear()
 
 
 def test_export_csv(driver: CustomDriver, html_doc):
@@ -204,11 +204,11 @@ def test_export_csv(driver: CustomDriver, html_doc):
     )
     assert content == expected
     csv_file.unlink()
-    register._clear()
+    register.clear()
 
 
 def test_selection_with_cache_check_and_uncheck(driver: CustomDriver, df):
-    register._init_grid("cached_sel")
+    register.add_grid("cached_sel")
     event = SimpleNamespace(new='{0: "CCC(C)CC"}')
     register.selection_updated("cached_sel", event)
     grid = get_grid(df, name="cached_sel", cache_selection=True)
@@ -219,7 +219,7 @@ def test_selection_with_cache_check_and_uncheck(driver: CustomDriver, df):
     assert sel == {0: "CCC(C)CC"}
     empty_sel = driver.click_checkbox(is_empty=True)
     assert empty_sel is True
-    register._clear()
+    register.clear()
 
 
 def test_selection_check_uncheck_invert(driver: CustomDriver, html_doc):
@@ -243,7 +243,7 @@ def test_selection_check_uncheck_invert(driver: CustomDriver, html_doc):
     driver.grid_action("invert")
     sel = driver.wait_for_selection(is_empty=False)
     assert len(sel) == 29
-    register._clear()
+    register.clear()
 
 
 @pytest.mark.parametrize("prerender", [True, False])

tests/test_marimo_integration.py

@@ -1,18 +1,19 @@
 import sys
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pandas as pd
 import pytest
 
 from mols2grid import MolGrid
+from mols2grid.select import register
 from mols2grid.utils import is_running_within_marimo
 from mols2grid.widget import MolGridWidget
 
 
 @pytest.fixture
-def mock_marimo_module():
-    with patch.dict(sys.modules, {"marimo": MagicMock()}):
-        yield
+def mock_marimo_module(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(sys, "modules", {**sys.modules, "marimo": MagicMock()})
 
 
 @pytest.mark.usefixtures("mock_marimo_module")
@@ -21,11 +22,9 @@ def test_is_running_within_marimo_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
+    if "marimo" in sys.modules:
+        assert isinstance(sys.modules["marimo"], MagicMock)
+    assert is_running_within_marimo() is False
 
 
 @pytest.fixture
@@ -52,13 +51,6 @@ def test_display_in_marimo(grid_fixture):
         assert result == mock_anywidget.return_value
 
 
-@pytest.mark.usefixtures("mock_marimo_module")
-def test_get_marimo_selection_before_rendering_raises(grid_fixture):
-    _, mg = grid_fixture
-    with pytest.raises(RuntimeError, match="run the `display` method first"):
-        mg.get_marimo_selection()
-
-
 @pytest.mark.usefixtures("mock_marimo_module")
 def test_get_selection_state_inside_marimo(grid_fixture):
     _, mg = grid_fixture
@@ -71,57 +63,53 @@ def test_get_selection_state_inside_marimo(grid_fixture):
         "marimo.state", return_value=(mock_get_state, mock_set_state)
     ) as mock_state:
         # Call get_marimo_selection
-        state_getter = mg.get_marimo_selection()
+        state_getter, _ = register.link_marimo_state()
 
         # Check if marimo.state was called with empty list
-        mock_state.assert_called_once_with([])
-
-        # Check if _marimo_hooked is set
-        assert getattr(mg.widget, "_marimo_hooked", False) is True
+        mock_state.assert_called_once_with({})
 
         # Verify return value
         assert state_getter == mock_get_state
 
 
-def test_get_selection_state_outside_marimo(grid_fixture):
-    _, mg = grid_fixture
-
+def test_get_selection_state_outside_marimo():
     # Ensure marimo is not in sys.modules
     with patch.dict(sys.modules):
         if "marimo" in sys.modules:
             del sys.modules["marimo"]
 
         with pytest.raises(RuntimeError, match="only available in a marimo notebook"):
-            mg.get_marimo_selection()
+            register.link_marimo_state()
 
 
 @pytest.mark.usefixtures("mock_marimo_module")
-def test_selection_state_update_logic(grid_fixture):
+def test_selection_state_update_logic(grid_fixture, monkeypatch: pytest.MonkeyPatch):
     _, mg = grid_fixture
-    mg.render()
-
     mock_set_state = MagicMock()
     with (
         patch("marimo.state", return_value=(MagicMock(), mock_set_state)),
-        patch.object(mg.widget, "observe") as mock_observe,
+        patch.object(MolGridWidget, "observe") as mock_observe,
     ):
         # Inspect the observe call to capture the callback
-        mg.get_marimo_selection()
+        register.link_marimo_state()
+        mock_callback = MagicMock(wraps=register.CALLBACKS[-1])
+        monkeypatch.setattr(register, "CALLBACKS", [mock_callback])
 
+        mg.render()
         # Verify observe was called
         mock_observe.assert_called()
-        args, _ = mock_observe.call_args
-        callback = args[0]
 
         # Simulate event with valid selection
-        # The widget returns a string representation of a dict
-        new_selection = {1: "C", 2: "CC"}
-        event = {"new": str(new_selection)}
-
-        callback(event)
-        mock_set_state.assert_called_with([1, 2])
-
-        # Test invalid input (should pass silently)
-        mock_set_state.reset_mock()
-        callback({"new": "invalid json"})
-        mock_set_state.assert_not_called()
+        event_values = {1: "C", 2: "CC"}
+        event = SimpleNamespace(new=str(event_values))
+        register.selection_updated("default", event)
+
+        # check callback was called with expected selection
+        mock_callback.assert_called_with("default", event_values)
+
+        # check inner lambda works as expected: given a current state with
+        # mol 42 selected, and new state where only 1 and 2 are selected,
+        # 42 should disappear and only 1, 2 remain
+        lambda_setter = mock_set_state.call_args[0][0]
+        result = lambda_setter({"default": [42]})
+        assert result == {"default": [1, 2]}

tests/test_molgrid.py

@@ -238,7 +238,7 @@ def test_get_selection(small_df):
     assert register.current_selection == "grid"
     assert grid.get_selection().equals(small_df.head(1))
     assert other.get_selection().equals(small_df.head(0))  # empty dataframe
-    register._clear()
+    register.clear()
 
 
 def test_save(grid, tmp_path):