From 16518bf6fe9ca018bf8576398deb9c8bc27d6dc3 Mon Sep 17 00:00:00 2001
From: Iisakki Rotko <iisakki.rotko@gmail.com>
Date: Fri, 26 Apr 2024 11:56:48 +0200
Subject: [PATCH] feat: check for key mutation when removing children (#35)

---
 reacton/core.py      |  4 ++++
 reacton/core_test.py | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/reacton/core.py b/reacton/core.py
index 60b1f8a..ac54e1a 100644
--- a/reacton/core.py
+++ b/reacton/core.py
@@ -230,6 +230,7 @@ def __init__(self, component, args=None, kwargs=None):
         self._meta = {}
         # for debugging/testing only
         self._render_count = 0
+        self._key_frozen: bool = False
 
         rc = _get_render_context(required=False)
         if rc is not None and rc.container_adders:
@@ -269,6 +270,8 @@ def key(self, value: str):
 
         This can help render performance. See documentation for details.
         """
+        if self._key_frozen:
+            raise RuntimeError("Element keys should not be mutated after rendering")
         self._key = value
         return self
 
@@ -1588,6 +1591,7 @@ def _render(self, element: Element, default_key: str, parent_key: str):
         key = el._key
         if key is None:
             key = default_key
+        el._key_frozen = True
 
         logger.debug("Render: (%s,%s)  - %r", parent_key, key, element)
 
diff --git a/reacton/core_test.py b/reacton/core_test.py
index ee967bc..9acd128 100644
--- a/reacton/core_test.py
+++ b/reacton/core_test.py
@@ -3186,3 +3186,25 @@ def Test():
     assert vbox.children[0].description == "1"
     assert vbox.children[1].description == "2"
     rc.close()
+
+
+def test_key_mutate_protection(Container):
+    button = ButtonComponentFunction(description="Hi")
+    set_state = lambda value: None  # noqa
+
+    @react.component
+    def Test():
+        nonlocal set_state
+        state, set_state = react.use_state(0)
+        if state == 0:
+            return Container(children=[button])
+        else:
+            return Container(children=[button.key("i_should_not_be_mutated")]).key("bla")
+
+    box, rc = react.render(Test(), handle_error=False)
+    assert rc.find(widgets.Button).widget.description == "Hi"
+    assert button._key is None
+    with pytest.raises(RuntimeError, match="Element keys should not be mutated after rendering"):
+        set_state(1)
+    rc.render(w.HTML(value="recover").key("HTML"))
+    rc.close()