Can't get this to work from inside a function

[ghost]

If I run the example code in a notebook it works perfectly:

import time
from ipywidgets import Button
from jupyter_ui_poll import ui_events

# Set up simple GUI, button with on_click callback
# that sets ui_done=True and changes button text
ui_done = False
def on_click(btn):
    global ui_done
    ui_done = True
    btn.description = '👍'

btn = Button(description='Click Me')
btn.on_click(on_click)
display(btn)

# Wait for user to press the button
with ui_events() as poll:
    while ui_done is False:
        poll(10)          # React to UI events (upto 10 at a time)
        print('.', end='')
        time.sleep(0.1)
print('done')

However, if I put this inside a function, as follows:

import time
from ipywidgets import Button
from jupyter_ui_poll import ui_events

def test():
  # Set up simple GUI, button with on_click callback
  # that sets ui_done=True and changes button text
  ui_done = False
  def on_click(btn):
      global ui_done
      ui_done = True
      btn.description = '👍'

  btn = Button(description='Click Me')
  btn.on_click(on_click)
  display(btn)

  # Wait for user to press the button
  with ui_events() as poll:
      while ui_done is False:
          poll(10)          # React to UI events (upto 10 at a time)
          print('.', end='')
          time.sleep(0.1)
  print('done')

Then call the function test(), then the code no longer responds to the button press, it just hangs with the button there, the block never finishes executing, and 'done' never gets printed.

---

Update:

My apologies this was an obvious error. ui_done wasn't set as a global variable in the main function, so on_click was looking at a different variable. My bad!