Revise, test, & document JupyterViz options for drawing agent space

Correct linter rerors flagged by ruff check

Author: rlskoeser

mesa/experimental/jupyter_viz.py

@@ -20,9 +20,23 @@ def JupyterViz(
     measures=None,
     name="Mesa Model",
     agent_portrayal=None,
-    space_drawer=None,
+    space_drawer="default",
     play_interval=400,
 ):
+    """Initialize a component to visualize a model.
+    Args:
+        model_class: class of the model to instantiate
+        model_params: parameters for initializing the model
+        measures: list of callables or data attributes to plot
+        name: name for display
+        agent_portrayal: options for rendering agents (dictionary)
+        space_drawer: method to render the agent space for
+            the model; default implementation is :meth:`make_space`;
+            simulations with no space to visualize should
+            specify `space_drawer=False`
+        play_interval: play interval (default: 400)
+    """
+
     current_step, set_current_step = solara.use_state(0)
 
     solara.Markdown(name)
@@ -48,10 +62,14 @@ def make_model():
 
     with solara.GridFixed(columns=2):
         # 4. Space
-        if space_drawer is None:
+        if space_drawer == "default":
+            # draw with the default implementation
             make_space(model, agent_portrayal)
-        else:
+        elif space_drawer:
+            # if specified, draw agent space with an alternate renderer
             space_drawer(model, agent_portrayal)
+        # otherwise, do nothing (do not draw space)
+
         # 5. Plots
         for measure in measures:
             if callable(measure):

tests/test_jupyter_viz.py

@@ -1,7 +1,9 @@
 import unittest
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
-from mesa.experimental.jupyter_viz import make_user_input
+import solara
+
+from mesa.experimental.jupyter_viz import JupyterViz, make_user_input
 
 
 class TestMakeUserInput(unittest.TestCase):
@@ -46,3 +48,54 @@ def test_label_fallback(self, mock_solara):
         mock_solara.SliderInt.assert_called_with(
             name, value=value, min=None, max=None, step=None
         )
+
+
+class TestJupyterViz(unittest.TestCase):
+    @patch("mesa.experimental.jupyter_viz.make_space")
+    def test_call_space_drawer(self, mock_make_space):
+        mock_model_class = Mock()
+        agent_portrayal = {
+            "Shape": "circle",
+            "color": "gray",
+        }
+        # initialize with space drawer unspecified (use default)
+        # component must be rendered for code to run
+        solara.render(
+            JupyterViz(
+                model_class=mock_model_class,
+                model_params={},
+                agent_portrayal=agent_portrayal,
+            )
+        )
+        # should call default method with class instance and agent portrayal
+        mock_make_space.assert_called_with(
+            mock_model_class.return_value, agent_portrayal
+        )
+
+        # specify no space should be drawn; any false value should work
+        for falsy_value in [None, False, 0]:
+            mock_make_space.reset_mock()
+            solara.render(
+                JupyterViz(
+                    model_class=mock_model_class,
+                    model_params={},
+                    agent_portrayal=agent_portrayal,
+                    space_drawer=falsy_value,
+                )
+            )
+            # should call default method with class instance and agent portrayal
+            assert mock_make_space.call_count == 0
+
+        # specify a custom space method
+        altspace_drawer = Mock()
+        solara.render(
+            JupyterViz(
+                model_class=mock_model_class,
+                model_params={},
+                agent_portrayal=agent_portrayal,
+                space_drawer=altspace_drawer,
+            )
+        )
+        altspace_drawer.assert_called_with(
+            mock_model_class.return_value, agent_portrayal
+        )