Skip to content

Commit

Permalink
Model: Replace get_agents_of_type method with agents_by_type property
Browse files Browse the repository at this point in the history
  • Loading branch information
EwoutH committed Sep 2, 2024
1 parent 0a1b2a3 commit 0da4a01
Show file tree
Hide file tree
Showing 4 changed files with 13 additions and 20 deletions.
4 changes: 2 additions & 2 deletions benchmarks/WolfSheep/wolf_sheep.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,8 +227,8 @@ def __init__(
patch.move_to(cell)

def step(self):
self.agents_of_type(Sheep).shuffle(inplace=True).do("step")
self.agents_of_type(Wolf).shuffle(inplace=True).do("step")
self.agents_by_type[Sheep].shuffle(inplace=True).do("step")
self.agents_by_type[Wolf].shuffle(inplace=True).do("step")


if __name__ == "__main__":
Expand Down
4 changes: 2 additions & 2 deletions mesa/experimental/devs/examples/wolf_sheep.py
Original file line number Diff line number Diff line change
Expand Up @@ -231,8 +231,8 @@ def __init__(
self.grid.place_agent(patch, pos)

def step(self):
self.agents_of_type(Sheep).shuffle(inplace=True).do("step")
self.agents_of_type(Wolf).shuffle(inplace=True).do("step")
self.agents_by_type[Sheep].shuffle(inplace=True).do("step")
self.agents_by_type[Wolf].shuffle(inplace=True).do("step")


if __name__ == "__main__":
Expand Down
23 changes: 8 additions & 15 deletions mesa/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,33 +106,26 @@ def agent_types(self) -> list[type]:
"""Return a list of all unique agent types registered with the model."""
return list(self._agents_by_type.keys())

def agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
"""Retrieves an AgentSet containing all agents of the specified type.
Args:
agenttype: The type of agent to retrieve.
Raises:
KeyError: If agenttype does not exist
"""
return self._agents_by_type[agenttype]
@property
def agents_by_type(self) -> dict[type[Agent], AgentSet]:
"""A dictionary where the keys are agent types and the values are the corresponding AgentSets."""
return self._agents_by_type

def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
"""Deprecated: Retrieves an AgentSet containing all agents of the specified type."""
warnings.warn(

Check warning on line 116 in mesa/model.py

View check run for this annotation

Codecov / codecov/patch

mesa/model.py#L116

Added line #L116 was not covered by tests
"get_agents_of_type is deprecated, please use agents_of_type instead.",
f"Model.get_agents_of_type() is deprecated, please replace get_agents_of_type({agenttype})"
f"with the property agents_by_type[{agenttype}].",
DeprecationWarning,
stacklevel=2,
)
return self.agents_of_type(agenttype)
return self.agents_by_type[agenttype]

Check warning on line 122 in mesa/model.py

View check run for this annotation

Codecov / codecov/patch

mesa/model.py#L122

Added line #L122 was not covered by tests

def _setup_agent_registration(self):
"""helper method to initialize the agent registration datastructures"""
self._agents = {} # the hard references to all agents in the model
self._agents_by_type: dict[
type, AgentSet
type[Agent], AgentSet
] = {} # a dict with an agentset for each class of agents
self._all_agents = AgentSet([], self) # an agenset with all agents

Expand Down
2 changes: 1 addition & 1 deletion tests/test_time.py
Original file line number Diff line number Diff line change
Expand Up @@ -320,7 +320,7 @@ def test_random_activation_counts(self):
agent_types = model.agent_types
for agent_type in agent_types:
assert model.schedule.get_type_count(agent_type) == len(
model.agents_of_type(agent_type)
model.agents_by_type[agent_type]
)

# def test_add_non_unique_ids(self):
Expand Down

0 comments on commit 0da4a01

Please sign in to comment.