Repository: projectmesa/mesa
Pull request #1960 (2024-01-13 11:18:00)

Bugfix in agentset to handle addition and removal correctly
Fix for #1959 and additional unittests that test specifically on this




codecov[bot]: ## [Codecov](https://app.codecov.io/gh/projectmesa/mesa/pull/1960?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`36da145`)](https://app.codecov.io/gh/projectmesa/mesa/commit/36da14542fa09e53497bb46a4e59bf757e29dcb0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 79.45% compared to head [(`4a14cda`)](https://app.codecov.io/gh/projectmesa/mesa/pull/1960?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 79.45%.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #1960   +/-   ##
=======================================
  Coverage   79.45%   79.45%           
=======================================
  Files          15       15           
  Lines        1285     1285           
  Branches      285      285           
=======================================
  Hits         1021     1021           
  Misses        225      225           
  Partials       39       39           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/projectmesa/mesa/pull/1960?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).

Corvince: pre-commit.ci autofix

Patch:
diff --git a/mesa/agent.py b/mesa/agent.py
index 0b204c22..109a9dcd 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -251,7 +251,12 @@ class AgentSet(MutableSet, Sequence):
         Returns:
             AgentSet | list[Any]: The results of the method calls if return_results is True, otherwise the AgentSet itself.
         """
-        res = [getattr(agent, method_name)(*args, **kwargs) for agent in self._agents]
+        # we iterate over the actual weakref keys and check if weakref is alive before calling the method
+        res = [
+            getattr(agentref(), method_name)(*args, **kwargs)
+            for agentref in self._agents.keyrefs()
+            if agentref()
+        ]
 
         return res if return_results else self
 


Test patch:
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 3cc8f9ad..5861038d 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -11,6 +11,26 @@ class TestAgent(Agent):
         return self.unique_id
 
 
+class TestAgentDo(Agent):
+    def __init__(
+        self,
+        unique_id,
+        model,
+    ):
+        super().__init__(unique_id, model)
+        self.agent_set = None
+
+    def get_unique_identifier(self):
+        return self.unique_id
+
+    def do_add(self):
+        agent = TestAgentDo(self.model.next_id(), self.model)
+        self.agent_set.add(agent)
+
+    def do_remove(self):
+        self.agent_set.remove(self)
+
+
 def test_agent_removal():
     model = Model()
     agent = TestAgent(model.next_id(), model)
@@ -164,6 +184,31 @@ def test_agentset_do_method():
     with pytest.raises(AttributeError):
         agentset.do("non_existing_method")
 
+    # tests for addition and removal in do
+    # do iterates, so no error should be raised to change size while iterating
+    # related to issue #1595
+
+    # setup
+    n = 10
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(n)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    agentset.do("do_add")
+    assert len(agentset) == 2 * n
+
+    # setup
+    model = Model()
+    agents = [TestAgentDo(model.next_id(), model) for _ in range(10)]
+    agentset = AgentSet(agents, model)
+    for agent in agents:
+        agent.agent_set = agentset
+
+    agentset.do("do_remove")
+    assert len(agentset) == 0
+
 
 def test_agentset_get_attribute():
     model = Model()