AgentSet: Add set method

mesa/agent.py

@@ -316,6 +316,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

@@ -270,6 +270,30 @@ 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)]