AgentSet: Add set method (#2254)

Adds a `set` method to the `AgentSet` class that allows setting a specified attribute of all agents within the set to a given value. This method provides a streamlined way to update agent attributes in bulk, without having to resort to `lambda` or other callable functions.

The motivation behind this feature is to simplify and expedite the process of modifying attributes across all agents in an `AgentSet`. Previously, this required iterating manually over the agents or using a `lambda` or other callable function. This new method makes the operation more elegant.

The `set` method was added to the `AgentSet` class. It iterates through all agents in the set and applies the `setattr` function to assign the specified value to the desired attribute. The method operates in place, modifying the existing agents directly.

The (now updated) AgentSet itself is returned, which allows for chaining.

You can both set existing Agent variables or set new ones.

mesa/agent.py

@@ -285,6 +285,21 @@ def get(self, attr_names: str | list[str]) -> list[Any]:
                 for agent in self._agents
             ]
 
+    def set(self, attr_name: str, value: Any) -> AgentSet:
+        """
+        Set a specified attribute to a given value for all agents in the AgentSet.
+
+        Args:
+            attr_name (str): The name of the attribute to set.
+            value (Any): The value to set the attribute to.
+
+        Returns:
+            AgentSet: The AgentSet instance itself, after setting the attribute.
+        """
+        for agent in self:
+            setattr(agent, attr_name, value)
+        return self
+
     def __getitem__(self, item: int | slice) -> Agent:
         """
         Retrieve an agent or a slice of agents from the AgentSet.

tests/test_agent.py

@@ -276,6 +276,59 @@ def remove_function(agent):
     assert len(agentset) == 0
 
 
+def test_agentset_set_method():
+    # Initialize the model and agents with and without existing attributes
+    class TestAgentWithAttribute(Agent):
+        def __init__(self, unique_id, model, age=None):
+            super().__init__(unique_id, model)
+            self.age = age
+
+    model = Model()
+    agents = [TestAgentWithAttribute(model.next_id(), model, age=i) for i in range(5)]
+    agentset = AgentSet(agents, model)
+
+    # Set a new attribute "health" and an existing attribute "age" for all agents
+    agentset.set("health", 100).set("age", 50).set("status", "active")
+
+    # Check if all agents have the "health", "age", and "status" attributes correctly set
+    for agent in agentset:
+        assert hasattr(agent, "health")
+        assert agent.health == 100
+        assert hasattr(agent, "age")
+        assert agent.age == 50
+        assert hasattr(agent, "status")
+        assert agent.status == "active"
+
+
+def test_agentset_map_str():
+    model = Model()
+    agents = [TestAgent(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+
+    with pytest.raises(AttributeError):
+        agentset.do("non_existing_method")
+
+    results = agentset.map("get_unique_identifier")
+    assert all(i == entry for i, entry in zip(results, range(1, 11)))
+
+
+def test_agentset_map_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.map(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
+
+    results = agentset.map(lambda agent: agent.unique_id)
+    assert all(i == entry for i, entry in zip(results, range(1, 11)))
+
+
 def test_agentset_get_attribute():
     model = Model()
     agents = [TestAgent(model.next_id(), model) for _ in range(10)]