Fix user_key and wid conflation bugs (#3827)

lib/streamlit/state/session_state.py

@@ -345,11 +345,11 @@ def keys(self) -> Set[str]:
         }
         return old_keys | new_widget_keys | new_session_state_keys
 
-    def is_new_state_value(self, key: str) -> bool:
-        return key in self._new_session_state
+    def is_new_state_value(self, user_key: str) -> bool:
+        return user_key in self._new_session_state
 
-    def is_new_widget_value(self, key: str) -> bool:
-        return key in self._new_widget_state
+    def is_new_widget_value(self, widget_id: str) -> bool:
+        return widget_id in self._new_widget_state
 
     def __iter__(self) -> Iterator[Any]:
         return iter(self.keys())
@@ -410,22 +410,23 @@ def _getitem(self, widget_id: Optional[str], user_key: Optional[str]) -> Any:
 
         raise KeyError
 
-    def __setitem__(self, key: str, value: Any) -> None:
+    def __setitem__(self, user_key: str, value: Any) -> None:
         from streamlit.report_thread import get_report_ctx
 
         ctx = get_report_ctx()
 
         if ctx is not None:
+            widget_id = self._key_id_mapping.get(user_key, None)
             widget_ids = ctx.widget_ids_this_run.items()
             form_ids = ctx.form_ids_this_run.items()
 
-            if key in widget_ids or key in form_ids:
+            if widget_id in widget_ids or user_key in form_ids:
                 raise StreamlitAPIException(
-                    f"`st.session_state.{key}` cannot be modified after the widget"
-                    f" with key `{key}` is instantiated."
+                    f"`st.session_state.{user_key}` cannot be modified after the widget"
+                    f" with key `{user_key}` is instantiated."
                 )
 
-        self._new_session_state[key] = value
+        self._new_session_state[user_key] = value
 
     def __delitem__(self, key: str) -> None:
         if key in INTERNAL_STATE_ATTRS:
@@ -524,25 +525,30 @@ def set_metadata(self, widget_metadata: WidgetMetadata) -> None:
         self._new_widget_state.widget_metadata[widget_id] = widget_metadata
 
     def maybe_set_new_widget_value(
-        self, widget_id: str, key: Optional[str] = None
+        self, widget_id: str, user_key: Optional[str] = None
     ) -> None:
         """Add the value of a new widget to session state."""
         widget_metadata = self._new_widget_state.widget_metadata[widget_id]
         deserializer = widget_metadata.deserializer
         initial_widget_value = deepcopy(deserializer(None, widget_metadata.id))
 
-        if widget_id not in self and (key is None or key not in self):
+        if widget_id not in self and (user_key is None or user_key not in self):
             # This is the first time this widget is being registered, so we save
             # its value in widget state.
             self._new_widget_state.set_from_value(widget_id, initial_widget_value)
 
-    def should_set_frontend_state_value(self, widget_id: str) -> bool:
+    def should_set_frontend_state_value(
+        self, widget_id: str, user_key: Optional[str]
+    ) -> bool:
         """Keep widget_state and session_state in sync when a widget is registered.
 
         This method returns whether the frontend needs to be updated with the
         new value of this widget.
         """
-        return self.is_new_state_value(widget_id) and not self.is_new_widget_value(
+        if user_key is None:
+            return False
+
+        return self.is_new_state_value(user_key) and not self.is_new_widget_value(
             widget_id
         )
 
@@ -563,8 +569,8 @@ def _get_widget_id(self, k: str) -> str:
         """
         return self._key_id_mapping.get(k, k)
 
-    def set_key_widget_mapping(self, widget_id: str, k: str) -> None:
-        self._key_id_mapping[k] = widget_id
+    def set_key_widget_mapping(self, widget_id: str, user_key: str) -> None:
+        self._key_id_mapping[user_key] = widget_id
 
     def copy(self):
         return deepcopy(self)

lib/streamlit/state/widgets.py

@@ -168,7 +168,7 @@ def register_widget(
         session_state.set_keyed_widget(metadata, widget_id, user_key)
     else:
         session_state.set_unkeyed_widget(metadata, widget_id)
-    value_changed = session_state.should_set_frontend_state_value(widget_id)
+    value_changed = session_state.should_set_frontend_state_value(widget_id, user_key)
 
     val = session_state.get_value_for_registration(widget_id)
 

lib/tests/streamlit/state/session_state_test.py

@@ -498,6 +498,7 @@ def test_setitem_disallows_setting_created_widget(self):
 
         with patch("streamlit.report_thread.get_report_ctx", return_value=mock_ctx):
             with pytest.raises(StreamlitAPIException) as e:
+                self.session_state._key_id_mapping = {"widget_id": "widget_id"}
                 self.session_state["widget_id"] = "blah"
             assert "`st.session_state.widget_id` cannot be modified" in str(e.value)
 
@@ -588,7 +589,8 @@ def test_should_set_frontend_state_value_new_widget(self):
         )
         assert (
             self.session_state.should_set_frontend_state_value(
-                f"{GENERATED_WIDGET_KEY_PREFIX}-0-widget_id_1"
+                f"{GENERATED_WIDGET_KEY_PREFIX}-0-widget_id_1",
+                "widget_id_1",
             )
             == False
         )