Multiselection refactoring

pages/0_👀_Lookup.py

@@ -9,7 +9,8 @@
     reduce_multiselect_font_size,
 )
 from st_widgets.console import session_console
-from st_widgets.stateful.data_editor import _update_data, stateful_data_editor
+from st_widgets.stateful import stateful_data_editor
+from st_widgets.stateful.data_editor import _update_data
 
 logging = get_logger(__name__)
 session = st.session_state
@@ -37,26 +38,17 @@ def load_dimensions(metabase2datasets: pd.DataFrame) -> pd.Series:
     codes = load_dimensions_and_codes(meta)
     codes = codes.reset_index().assign(selected=False)
 
-    with st.form("_selected_codes_form"):
-        key = "_selected_codes"
-
-        if f"{key}_data" not in session:
-            session[f"{key}_data"] = codes
-
-        st.data_editor(
-            session[f"{key}_data"],
-            disabled=["code", "dimension", "description"],
-            use_container_width=True,
-            key=key,
-            on_change=_update_data(session, key),
-        )
-
-        submitted = st.form_submit_button("Add")
-        if submitted:
-            selected_codes = session[f"{key}_data"]
-            selected_codes_mask = selected_codes["selected"].values
-        else:
-            selected_codes_mask = codes["selected"].values
+    selected_codes = stateful_data_editor(
+        codes,
+        disabled=["code", "dimension", "description"],
+        use_container_width=True,
+        key="_selected_codes",
+        multiedit=True,
+    )
+    if selected_codes is not None:
+        selected_codes_mask = selected_codes["selected"].values
+    else:
+        selected_codes_mask = codes["selected"].values
 
     # TODO  Retrieve selected dimension after page change
     st.markdown("Selected dimension overview:")

st_widgets/stateful/data_editor.py

@@ -28,18 +28,40 @@ def stateful_data_editor(
     position: DeltaGenerator = st._main,
     session: MutableMapping[Key, Any] = st.session_state,
     on_change: Optional[WidgetCallback] = None,
+    multiedit: bool = False,
     **kwargs,
 ):
     """
     A stateful data editor that preserves modification.
+    Can be configured to accept multiple editing before reload (performed on button click).
     """
     if f"{key}_data" not in session:
         session[f"{key}_data"] = data
 
-    position.data_editor(
-        data=session[f"{key}_data"],
-        key=key,
-        on_change=_on_change_factory(partial(_update_data, session, key))(on_change),
-        **kwargs,
-    )
-    return session[f"{key}_data"]
+    if multiedit:
+        with position.form(f"{key}_form"):
+            position.data_editor(
+                data=session[f"{key}_data"],
+                key=key,
+                # TODO For unknown reasons, use `_on_change_factory` here cause:
+                # `StreamlitAPIException: With forms, callbacks can only be defined on the st.form_submit_button.
+                # Defining callbacks on other widgets inside a form is not allowed.`
+                # Anyway, `_update_data` IS a callback and it is working.
+                on_change=_update_data(session, key),
+                **kwargs,
+            )
+
+            submitted = position.form_submit_button("Edit")
+            if submitted:
+                return session[f"{key}_data"]
+
+    else:
+        position.data_editor(
+            data=session[f"{key}_data"],
+            key=key,
+            on_change=_on_change_factory(partial(_update_data, session, key))(
+                on_change
+            ),
+            **kwargs,
+        )
+        return session[f"{key}_data"]