Introduce AgentSet class

mesa/agent.py

@@ -7,10 +7,11 @@
 # Remove this __future__ import once the oldest supported Python is 3.10
 from __future__ import annotations
 
+import weakref
 from random import Random
 
 # mypy
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Iterable
 
 if TYPE_CHECKING:
     # We ensure that these are not imported during runtime to prevent cyclic
@@ -41,12 +42,15 @@ def __init__(self, unique_id: int, model: Model) -> None:
         self.model = model
         self.pos: Position | None = None
 
-        # Register the agent with the model using defaultdict
-        self.model.agents[type(self)][self] = None
+        # Directly register the agent with the model
+        if type(self) not in self.model._agents:
+            self.model._agents[type(self)] = set()
+        self.model._agents[type(self)].add(self)
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
-        self.model.agents[type(self)].pop(self)
+        if type(self) in self.model._agents:
+            self.model._agents[type(self)].discard(self)
 
     def step(self) -> None:
         """A single step of the agent."""
@@ -57,3 +61,90 @@ def advance(self) -> None:
     @property
     def random(self) -> Random:
         return self.model.random
+
+
+class AgentSet:
+    """Ordered set of agents"""
+
+    def __init__(self, agents: Iterable[Agent], model: Model):
+        self._agents = weakref.WeakKeyDictionary()
+        self._indices = []
+        for agent in agents:
+            self._agents[agent] = None
+            self._indices.append(weakref.ref(agent))
+
+        self.model = model
+
+    def __len__(self) -> int:
+        return len(self._agents)
+
+    def __iter__(self) -> Iterator[Agent]:
+        return self._agents.keys()
+
+    def __contains__(self, agent: Agent) -> bool:
+        """Check if an agent is in the AgentSet."""
+        return agent in self._agents
+
+    def select(self, filter_func: Callable[[Agent], bool] | None = None, inplace: bool = False):
+        if filter_func is None:
+            return self if inplace else AgentSet(list(self._agents.keys()), self.model)
+
+        agents = [agent for agent in self._agents.keys() if filter_func(agent)]
+
+        if inplace:
+            self._reorder(agents)
+            return self
+        else:
+            return AgentSet(
+                agents,
+                self.model
+            )
+
+    def shuffle(self, inplace: bool = False):
+        shuffled_agents = list(self._agents.keys())
+        self.model.random.shuffle(shuffled_agents)
+
+        if inplace:
+            self._reorder(shuffled_agents)
+            return self
+        else:
+            return AgentSet(shuffled_agents, self.model)
+
+    def sort(self, key: Callable[[Agent], Any], reverse: bool = False, inplace: bool = False):
+        sorted_agents = sorted(self._agents.keys(), key=key, reverse=reverse)
+
+        if inplace:
+            self._reorder(sorted_agents)
+            return self
+        else:
+            return AgentSet(sorted_agents, self.model)
+
+    def _reorder(self, agents: Iterable[Agent]):
+        _agents = weakref.WeakKeyDictionary()
+        _indices = []
+        for agent in agents:
+            _agents[agent] = None
+            _indices.append(weakref.ref(agent))
+        self._agents = _agents
+        self._indices = _indices
+
+    def do_each(self, method_name: str, *args, **kwargs) -> list[Any]:
+        """invoke method on each agent"""
+        return [getattr(agent, method_name)(*args, **kwargs) for agent in self._agents]
+
+    def get_each(self, attr_name: str) -> list[Any]:
+        """get attribute value on each agent"""
+        return [getattr(agent, attr_name) for agent in self._agents]
+
+    def __getitem__(self, item: int) -> Agent:
+        # TODO::
+        # TBD:: it is a bit tricky to make this work
+        # part of the problem is that there is no weakreflist
+        agent = self._indices[item]()
+        if agent is None:
+            # the agent has been garbage collected
+            return None
+        else:
+            return self._agents[agent]
+
+    # Additional methods like union, intersection, difference, etc., can be added as needed.

mesa/model.py

@@ -13,6 +13,7 @@
 # mypy
 from typing import Any
 
+from mesa.agent import AgentSet
 from mesa.datacollection import DataCollector
 
 
@@ -51,12 +52,22 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.running = True
         self.schedule = None
         self.current_id = 0
-        self.agents: defaultdict[type, dict] = defaultdict(dict)
+        self._agents: defaultdict[type, dict] = defaultdict(dict)
+
+    @property
+    def agents(self) -> AgentSet:
+        all_agents = set()
+        for agent_type in self._agents:
+            all_agents.update(self._agents[agent_type])
+        return AgentSet(all_agents, self)
 
     @property
     def agent_types(self) -> list:
         """Return a list of different agent types."""
-        return list(self.agents.keys())
+        return list(self._agents.keys())
+
+    def select_agents(self, *args, **kwargs) -> AgentSet:
+        return self.agents.select(*args, **kwargs)
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached. Overload as

tests/test_agent.py

@@ -9,8 +9,8 @@ class TestAgent(Agent):
     model = Model()
     agent = TestAgent(model.next_id(), model)
     # Check if the agent is added
-    assert agent in model.agents[type(agent)]
+    assert agent in model.agents
 
     agent.remove()
     # Check if the agent is removed
-    assert agent not in model.agents[type(agent)]
+    assert agent not in model.agents

tests/test_model.py

@@ -49,5 +49,5 @@ class TestAgent(Agent):
 
     model = Model()
     test_agent = TestAgent(model.next_id(), model)
-    assert test_agent in model.agents[type(test_agent)]
+    assert test_agent in model.agents
     assert type(test_agent) in model.agent_types