diff --git a/mesa/agent.py b/mesa/agent.py index 18454164e5a..200dab6f093 100644 --- a/mesa/agent.py +++ b/mesa/agent.py @@ -67,7 +67,13 @@ def __init__(self, model: Model, *args, **kwargs) -> None: self.model.register_agent(self) def remove(self) -> None: - """Remove and delete the agent from the model.""" + """Remove and delete the agent from the model. + + Notes: + If you need to do additional cleanup when removing an agent by for example removing + it from a space, consider extending this method in your own agent class. + + """ with contextlib.suppress(KeyError): self.model.deregister_agent(self) diff --git a/mesa/model.py b/mesa/model.py index ac59cc1f1f9..f9c0e2e4869 100644 --- a/mesa/model.py +++ b/mesa/model.py @@ -276,3 +276,16 @@ def initialize_data_collector( ) # Collect data for the first time during initialization. self.datacollector.collect(self) + + def remove_all_agents(self): + """Remove all agents from the model. + + Notes: + This method calls agent.remove for all agents in the model. If you need to remove agents from + e.g., a SingleGrid, you can either explicitly implement your own agent.remove method or clean this up + near where you are calling this method. + + """ + # we need to wrap keys in a list to avoid a RunTimeError: dictionary changed size during iteration + for agent in list(self._agents.keys()): + agent.remove() diff --git a/tests/test_model.py b/tests/test_model.py index 7c343e46357..ff1fe297b6a 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -93,3 +93,18 @@ class Sheep(Agent): assert model.agents_by_type[Wolf] == AgentSet([wolf], model) assert model.agents_by_type[Sheep] == AgentSet([sheep], model) assert len(model.agents_by_type) == 2 + + +def test_agent_remove(): + """Test removing all agents from the model.""" + + class TestAgent(Agent): + pass + + model = Model() + for _ in range(100): + TestAgent(model) + assert len(model.agents) == 100 + + model.remove_all_agents() + assert len(model.agents) == 0