solara: Implement decorator to simplify custom space drawer

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

mesa/experimental/jupyter_viz.py

@@ -105,15 +105,22 @@ def _draw_network_grid(viz, space_ax):
     )
 
 
-def make_space(viz):
-    space_fig = Figure()
-    space_ax = space_fig.subplots()
+def prepare_matplotlib_space(drawer):
+    def wrapped_drawer(viz):
+        space_fig = Figure()
+        space_ax = space_fig.subplots()
+        drawer(viz, space_fig, space_ax)
+        space_ax.set_axis_off()
+        solara.FigureMatplotlib(space_fig, dependencies=[viz.model, viz.df])
+    return wrapped_drawer
+
+
+@prepare_matplotlib_space
+def make_space(viz, space_fig, space_ax):
     if isinstance(viz.model.grid, mesa.space.NetworkGrid):
         _draw_network_grid(viz, space_ax)
     else:
         space_ax.scatter(**viz.portray(viz.model.grid))
-    space_ax.set_axis_off()
-    solara.FigureMatplotlib(space_fig, dependencies=[viz.model, viz.df])
 
 
 def make_plot(viz, measure):