Allow AgentSet.do() to take Callable function

mesa/agent.py

@@ -227,26 +227,37 @@ def _update(self, agents: Iterable[Agent]):
         return self
 
     def do(
-        self, method_name: str, *args, return_results: bool = False, **kwargs
+        self, method: str | Callable, *args, return_results: bool = False, **kwargs
     ) -> AgentSet | list[Any]:
         """
-        Invoke a method on each agent in the AgentSet.
+        Invoke a method or function on each agent in the AgentSet.
 
         Args:
-            method_name (str): The name of the method to call on each agent.
+            method (str, callable): the callable to do on each agents
+
+                                        * in case of str, the name of the method to call on each agent.
+                                        * in case of callable, the function to be called with each agent as first argument
+
             return_results (bool, optional): If True, returns the results of the method calls; otherwise, returns the AgentSet itself. Defaults to False, so you can chain method calls.
-            *args: Variable length argument list passed to the method being called.
-            **kwargs: Arbitrary keyword arguments passed to the method being called.
+            *args: Variable length argument list passed to the callable being called.
+            **kwargs: Arbitrary keyword arguments passed to the callable being called.
 
         Returns:
-            AgentSet | list[Any]: The results of the method calls if return_results is True, otherwise the AgentSet itself.
+            AgentSet | list[Any]: The results of the callable calls if return_results is True, otherwise the AgentSet itself.
         """
         # we iterate over the actual weakref keys and check if weakref is alive before calling the method
-        res = [
-            getattr(agent, method_name)(*args, **kwargs)
-            for agentref in self._agents.keyrefs()
-            if (agent := agentref()) is not None
-        ]
+        if isinstance(method, str):
+            res = [
+                getattr(agent, method)(*args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
+        else:
+            res = [
+                method(agent, *args, **kwargs)
+                for agentref in self._agents.keyrefs()
+                if (agent := agentref()) is not None
+            ]
 
         return res if return_results else self
 

tests/test_agent.py

@@ -176,7 +176,7 @@ def test_agentset_get_item():
         _ = agentset[20]
 
 
-def test_agentset_do_method():
+def test_agentset_do_str():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]
     agentset = AgentSet(agents, model)
@@ -210,6 +210,72 @@ def test_agentset_do_method():
     assert len(agentset) == 0
 
 
+def test_agentset_do_callable():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    # Test callable with non-existent function
+    with pytest.raises(AttributeError):
+        agentset.do(lambda agent: agent.non_existing_method())
+
+    # tests for addition and removal in do using callables
+    # do iterates, so no error should be raised to change size while iterating
+    # related to issue #1595
+
+    # setup for lambda function tests
+    n = 10
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(n)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Lambda for addition
+    agentset.do(lambda agent: agent.do_add())
+    assert len(agentset) == 2 * n
+
+    # setup again for lambda function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Lambda for removal
+    agentset.do(lambda agent: agent.do_remove())
+    assert len(agentset) == 0
+
+    # setup for actual function tests
+    def add_function(agent):
+        agent.do_add()
+
+    def remove_function(agent):
+        agent.do_remove()
+
+    # setup again for actual function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(n)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Actual function for addition
+    agentset.do(add_function)
+    assert len(agentset) == 2 * n
+
+    # setup again for actual function tests
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    # Actual function for removal
+    agentset.do(remove_function)
+    assert len(agentset) == 0
+
+
 def test_agentset_get_attribute():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]