From f860a060ed9ba7048004add48f044b1169994dae Mon Sep 17 00:00:00 2001
From: rht <rhtbot@protonmail.com>
Date: Tue, 14 Nov 2023 20:01:30 -0500
Subject: [PATCH] refactor: Define current_step as reactive

This should reduce the number of arguments that needs to be passed to
the components.
---
 mesa/experimental/jupyter_viz.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
index c613131ba83..062c0dc2005 100644
--- a/mesa/experimental/jupyter_viz.py
+++ b/mesa/experimental/jupyter_viz.py
@@ -38,7 +38,7 @@ def JupyterViz(
             specify `space_drawer=False`
         play_interval: play interval (default: 150)
     """
-    current_step, set_current_step = solara.use_state(0)
+    current_step = solara.use_reactive(0)
 
     # 1. Set up model parameters
     user_params, fixed_params = split_model_params(model_params)
@@ -49,7 +49,7 @@ def JupyterViz(
     # 2. Set up Model
     def make_model():
         model = model_class(**model_parameters)
-        set_current_step(0)
+        current_step.value = 0
         return model
 
     reset_counter = solara.use_reactive(0)
@@ -97,7 +97,7 @@ def render_in_jupyter():
         with solara.GridFixed(columns=2):
             UserInputs(user_params, on_change=handle_change_model_params)
             ModelController(
-                model, play_interval, current_step, set_current_step, reset_counter
+                model, play_interval, current_step, reset_counter
             )
             solara.Markdown(md_text=f"###Step - {current_step}")
 
@@ -134,7 +134,7 @@ def render_in_browser():
             with solara.Card("Controls", margin=1, elevation=2):
                 UserInputs(user_params, on_change=handle_change_model_params)
                 ModelController(
-                    model, play_interval, current_step, set_current_step, reset_counter
+                    model, play_interval, current_step, reset_counter
                 )
             with solara.Card("Progress", margin=1, elevation=2):
                 solara.Markdown(md_text=f"####Step - {current_step}")
@@ -161,7 +161,7 @@ def render_in_browser():
 
 @solara.component
 def ModelController(
-    model, play_interval, current_step, set_current_step, reset_counter
+    model, play_interval, current_step, reset_counter
 ):
     playing = solara.use_reactive(False)
     thread = solara.use_reactive(None)
@@ -172,8 +172,8 @@ def ModelController(
     previous_step = solara.use_reactive(0)
 
     def on_value_play(change):
-        if previous_step.value > current_step and current_step == 0:
-            # We add extra checks for current_step == 0, just to be sure.
+        if previous_step.value > current_step.value and current_step.value == 0:
+            # We add extra checks for current_step.value == 0, just to be sure.
             # We automatically stop the playing if a model is reset.
             playing.value = False
         elif model.running:
@@ -183,8 +183,8 @@ def on_value_play(change):
 
     def do_step():
         model.step()
-        previous_step.value = current_step
-        set_current_step(model.schedule.steps)
+        previous_step.value = current_step.value
+        current_step.value = model.schedule.steps
 
     def do_play():
         model.running = True
@@ -207,11 +207,11 @@ def do_reset():
         reset_counter.value += 1
 
     def do_set_playing(value):
-        if current_step == 0:
+        if current_step.value == 0:
             # This means the model has been recreated, and the step resets to
             # 0. We want to avoid triggering the playing.value = False in the
             # on_value_play function.
-            previous_step.value = current_step
+            previous_step.value = current_step.value
         playing.set(value)
 
     with solara.Row():