fixes weakref bug in shuffe_do

mesa/agent.py

@@ -294,15 +294,17 @@ def shuffle_do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
 
         It's a fast, optimized version of calling shuffle() followed by do().
         """
-        agents = list(self._agents.keys())
-        self.random.shuffle(agents)
+        weakrefs = list(self._agents.keyrefs())
+        self.random.shuffle(weakrefs)
 
         if isinstance(method, str):
-            for agent in agents:
-                getattr(agent, method)(*args, **kwargs)
+            for ref in weakrefs:
+                if (agent := ref()) is not None:
+                    getattr(agent, method)(*args, **kwargs)
         else:
-            for agent in agents:
-                method(agent, *args, **kwargs)
+            for ref in weakrefs:
+                if (agent := ref()) is not None:
+                    method(agent, *args, **kwargs)
 
         return self
 

tests/test_agent.py

@@ -472,6 +472,29 @@ def test_method(self):
         original_order != shuffled_order
     ), "The order should be different after shuffle_do"
 
+    class AgentWithRemove(Agent):
+        def __init__(self, model):
+            super().__init__(model)
+            self.is_alive = True
+
+        def remove(self):
+            super().remove()
+            self.is_alive = False
+
+        def step(self):
+            if not self.is_alive:
+                raise Exception
+
+            agent_to_remove = self.random.choice(self.model.agents)
+
+            if agent_to_remove is not self:
+                agent_to_remove.remove()
+
+    model = Model(seed=32)
+    for _ in range(100):
+        AgentWithRemove(model)
+    model.agents.shuffle_do("step")
+
 
 def test_agentset_get_attribute():
     """Test AgentSet.get for attributes."""