Skip to content

Commit de58ddb

Browse files
authored
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. ### Motive 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. ### Implementation 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.
1 parent 3e25d9c commit de58ddb

File tree

2 files changed

+39
-0
lines changed

2 files changed

+39
-0
lines changed

mesa/agent.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,21 @@ def get(self, attr_names: str | list[str]) -> list[Any]:
316316
for agent in self._agents
317317
]
318318

319+
def set(self, attr_name: str, value: Any) -> AgentSet:
320+
"""
321+
Set a specified attribute to a given value for all agents in the AgentSet.
322+
323+
Args:
324+
attr_name (str): The name of the attribute to set.
325+
value (Any): The value to set the attribute to.
326+
327+
Returns:
328+
AgentSet: The AgentSet instance itself, after setting the attribute.
329+
"""
330+
for agent in self:
331+
setattr(agent, attr_name, value)
332+
return self
333+
319334
def __getitem__(self, item: int | slice) -> Agent:
320335
"""
321336
Retrieve an agent or a slice of agents from the AgentSet.

tests/test_agent.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,30 @@ def remove_function(agent):
270270
assert len(agentset) == 0
271271

272272

273+
def test_agentset_set_method():
274+
# Initialize the model and agents with and without existing attributes
275+
class TestAgentWithAttribute(Agent):
276+
def __init__(self, unique_id, model, age=None):
277+
super().__init__(unique_id, model)
278+
self.age = age
279+
280+
model = Model()
281+
agents = [TestAgentWithAttribute(model.next_id(), model, age=i) for i in range(5)]
282+
agentset = AgentSet(agents, model)
283+
284+
# Set a new attribute "health" and an existing attribute "age" for all agents
285+
agentset.set("health", 100).set("age", 50).set("status", "active")
286+
287+
# Check if all agents have the "health", "age", and "status" attributes correctly set
288+
for agent in agentset:
289+
assert hasattr(agent, "health")
290+
assert agent.health == 100
291+
assert hasattr(agent, "age")
292+
assert agent.age == 50
293+
assert hasattr(agent, "status")
294+
assert agent.status == "active"
295+
296+
273297
def test_agentset_map_str():
274298
model = Model()
275299
agents = [TestAgent(model.next_id(), model) for _ in range(10)]

0 commit comments

Comments
 (0)