solara: Implement decorator to simplify custom space drawer

projectmesa · Sep 1, 2023 · 039033f · 039033f
1 parent 47637a7
commit 039033f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/mesa/experimental/__init__.py b/mesa/experimental/__init__.py
@@ -1 +1 @@
-from .jupyter_viz import JupyterViz, make_text  # noqa
+from .jupyter_viz import JupyterViz, make_text, prepare_matplotlib_space  # noqa
diff --git a/mesa/experimental/jupyter_viz.py b/mesa/experimental/jupyter_viz.py
@@ -182,7 +182,19 @@ def make_user_input(user_input, name, options):
         raise ValueError(f"{input_type} is not a supported input type")
 
 
-def make_space(model, agent_portrayal):
+def prepare_matplotlib_space(drawer):
+    def wrapped_drawer(model, agent_portrayal):
+        space_fig = Figure()
+        space_ax = space_fig.subplots()
+        drawer(model, agent_portrayal, space_fig, space_ax)
+        space_ax.set_axis_off()
+        solara.FigureMatplotlib(space_fig)
+
+    return wrapped_drawer
+
+
+@prepare_matplotlib_space
+def make_space(model, agent_portrayal, space_fig, space_ax):
     def portray(g):
         x = []
         y = []
@@ -211,14 +223,10 @@ def portray(g):
             out["c"] = c
         return out
 
-    space_fig = Figure()
-    space_ax = space_fig.subplots()
     if isinstance(model.grid, mesa.space.NetworkGrid):
         _draw_network_grid(model, space_ax, agent_portrayal)
     else:
         space_ax.scatter(**portray(model.grid))
-    space_ax.set_axis_off()
-    solara.FigureMatplotlib(space_fig)
 
 
 def _draw_network_grid(model, space_ax, agent_portrayal):
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		from .jupyter_viz import JupyterViz, make_text # noqa
		from .jupyter_viz import JupyterViz, make_text, prepare_matplotlib_space # noqa