regression bug (0.5.2->0.6.0): batch wrapped Switch looses it's on_change handler

[mrdobalina2k]

Describe the bug

In marimo 0.5.2, the code below works fine and prints out 1, 2, 3, ... whenever the wrapped switch is clicked. In marimo 0.6.0, this is no longer the case. Only the first click makes the switch call it's on_change handler. I would expect the on_change handler to fire every time.

Environment

{
"marimo": "0.6.0",
"OS": "Windows",
"OS Version": "11",
"Processor": "Intel64 Family 6 Model 126 Stepping 5, GenuineIntel",
"Python Version": "3.12.2",
"Binaries": {
"Browser": "123.0.6312.123",
"Node": "v14.16.0"
},
"Requirements": {
"click": "8.1.7",
"importlib-resources": "missing",
"jedi": "0.19.1",
"markdown": "3.6",
"pymdown-extensions": "10.7.1",
"pygments": "2.17.2",
"tomlkit": "0.12.4",
"uvicorn": "0.29.0",
"starlette": "0.37.2",
"websocket": "missing",
"typing-extensions": "4.9.0",
"black": "24.3.0"
}
}

Code to reproduce

here's a minimal repro:

import marimo

__generated_with = "0.6.0"
app = marimo.App()


@app.cell
def __():
    import marimo as mo
    return mo,


@app.cell
def __(mo):
    get_state, set_state = mo.state(0)
    return get_state, set_state


@app.cell
def __(get_state, mo, set_state):
    def on_change(value):
        v = get_state()
        set_state(v+1)

    switch = mo.ui.switch(value=True, on_change=on_change)
    return on_change, switch


@app.cell
def __(switch):
    switch
    return


@app.cell
def __(mo, switch):
    def wrap_switch(switch):

        wrapped = mo.md("""
        wrapped switch

        {switch}""").batch(switch=switch)
        return wrapped
    wrapped_switch = wrap_switch(switch)
    wrapped_switch
    return wrap_switch, wrapped_switch


@app.cell
def __(get_state):
    get_state()
    return


if __name__ == "__main__":
    app.run()

[akshayka]

Thanks for reporting, and for the reproduction.

On >= 0.5.2, the handler is being called, but I'm also seeing a bug in which the batched switch UI element is resetting after being clicked. Do you also see that?

Screencast.from.05-21-2024.03.15.14.PM.webm

[akshayka]

@mrdobalina2k : Thanks for reporting, I have a fix out in #1439.

By the way, I figured out the flickering/resetting of the slider. It's because your on_change handler calls get_state(). To fix that, use the function API for set_state, like below:

import marimo

__generated_with = "0.6.0"
app = marimo.App()


@app.cell
def __():
    import marimo as mo
    return mo,


@app.cell
def __(mo):
    get_state, set_state = mo.state(0)
    return get_state, set_state


@app.cell
def __(mo, set_state):
    def on_change(value):
        # don't call `get_state` here!
        set_state(lambda v: v+1)

    switch = mo.ui.switch(value=True, on_change=on_change)
    return on_change, switch


@app.cell
def __(switch):
    switch
    return


@app.cell
def __(mo, switch):
    def wrap_switch(switch):

        wrapped = mo.md("""
        wrapped switch

        {switch}""").batch(switch=switch)
        return wrapped

    wrapped_switch = wrap_switch(switch)
    wrapped_switch
    return wrap_switch, wrapped_switch


@app.cell
def __(wrapped_switch):
    wrapped_switch._elements["switch"]._id
    return


@app.cell
def __(wrapped_switch):
    wrapped_switch.value
    return


@app.cell
def __():
    from marimo._runtime.context import get_context

    get_context().ui_element_registry._objects
    return get_context,


@app.cell
def __(get_state):
    get_state()
    return


if __name__ == "__main__":
    app.run()