feat: Implement auto-conversion of function to TextElement

examples/schelling/server.py

@@ -3,16 +3,11 @@
 from model import Schelling
 
 
-class HappyElement(mesa.visualization.TextElement):
+def get_happy_agents(model):
     """
     Display a text count of how many happy agents there are.
     """
-
-    def __init__(self):
-        pass
-
-    def render(self, model):
-        return "Happy agents: " + str(model.happy)
+    return f"Happy agents: {model.happy}"
 
 
 def schelling_draw(agent):
@@ -32,7 +27,6 @@ def schelling_draw(agent):
     return portrayal
 
 
-happy_element = HappyElement()
 canvas_element = mesa.visualization.CanvasGrid(schelling_draw, 20, 20, 500, 500)
 happy_chart = mesa.visualization.ChartModule([{"Label": "happy", "Color": "Black"}])
 
@@ -45,5 +39,8 @@ def schelling_draw(agent):
 }
 
 server = mesa.visualization.ModularServer(
-    Schelling, [canvas_element, happy_element, happy_chart], "Schelling", model_params
+    Schelling,
+    [canvas_element, get_happy_agents, happy_chart],
+    "Schelling",
+    model_params,
 )

examples/virus_on_network/virus_on_network/server.py

@@ -59,15 +59,14 @@ def get_agents(source, target):
 )
 
 
-class MyTextElement(mesa.visualization.TextElement):
-    def render(self, model):
-        ratio = model.resistant_susceptible_ratio()
-        ratio_text = "∞" if ratio is math.inf else f"{ratio:.2f}"
-        infected_text = str(number_infected(model))
-
-        return "Resistant/Susceptible Ratio: {}<br>Infected Remaining: {}".format(
-            ratio_text, infected_text
-        )
+def get_resistant_susceptible_ratio(model):
+    ratio = model.resistant_susceptible_ratio()
+    ratio_text = "&infin;" if ratio is math.inf else f"{ratio:.2f}"
+    infected_text = str(number_infected(model))
+
+    return "Resistant/Susceptible Ratio: {}<br>Infected Remaining: {}".format(
+        ratio_text, infected_text
+    )
 
 
 model_params = {
@@ -126,6 +125,9 @@ def render(self, model):
 }
 
 server = mesa.visualization.ModularServer(
-    VirusOnNetwork, [network, MyTextElement(), chart], "Virus Model", model_params
+    VirusOnNetwork,
+    [network, get_resistant_susceptible_ratio, chart],
+    "Virus Model",
+    model_params,
 )
 server.port = 8521

mesa/visualization/ModularVisualization.py

@@ -164,6 +164,15 @@ def render(self, model):
         return "<b>VisualizationElement goes here</b>."
 
 
+class TextElement(VisualizationElement):
+    """
+    Module for drawing live-updating text.
+    """
+
+    package_includes = ["TextModule.js"]
+    js_code = "elements.push(new TextModule());"
+
+
 # =============================================================================
 # Actual Tornado code starts here:
 
@@ -273,7 +282,9 @@ def __init__(
         if model_params is None:
             model_params = {}
         # Prep visualization elements:
-        self.visualization_elements = visualization_elements
+        self.visualization_elements = self._auto_convert_functions_to_TextElements(
+            visualization_elements
+        )
         self.package_js_includes = set()
         self.package_css_includes = set()
         self.local_js_includes = set()
@@ -357,3 +368,26 @@ def launch(self, port=None, open_browser=True):
     @staticmethod
     def _is_stylesheet(filename):
         return filename.lower().endswith(".css")
+
+    def _auto_convert_fn_to_TextElement(self, x):
+        """
+        Automatically convert a function to a TextElement object.
+        See https://github.com/projectmesa/mesa/issues/1233.
+        """
+
+        # Note: a class constructor is also a callable.
+        if not callable(x):
+            # i.e. not a function
+            return x
+
+        class MyTextElement(TextElement):
+            def render(self, model):
+                return x(model)
+
+        return MyTextElement()
+
+    def _auto_convert_functions_to_TextElements(self, visualization_elements):
+        out_elements = [
+            self._auto_convert_fn_to_TextElement(e) for e in visualization_elements
+        ]
+        return out_elements

mesa/visualization/modules/__init__.py

@@ -6,6 +6,9 @@
 from mesa.visualization.modules.ChartVisualization import ChartModule  # noqa
 from mesa.visualization.modules.PieChartVisualization import PieChartModule  # noqa
 from mesa.visualization.modules.BarChartVisualization import BarChartModule  # noqa
-from mesa.visualization.modules.TextVisualization import TextElement  # noqa
 from mesa.visualization.modules.HexGridVisualization import CanvasHexGrid  # noqa
 from mesa.visualization.modules.NetworkVisualization import NetworkModule  # noqa
+
+# Delete this line in the next major release, once the simpler namespace has
+# become widely adopted.
+from mesa.visualization.ModularVisualization import TextElement  # noqa