Merge pull request projectmesa#1001 from Casper-Smet/remove_agent

Public remove_agent function for NetworkGrid

mesa/space.py

@@ -950,6 +950,12 @@ def _remove_agent(self, agent: Agent, node_id: int) -> None:
 
         self.G.nodes[node_id]["agent"].remove(agent)
 
+    def remove_agent(self, agent: Agent) -> None:
+        """ Remove the agent from the network and set its pos variable to None. """
+        pos = agent.pos
+        self._remove_agent(agent, pos)
+        agent.pos = None
+
     def is_cell_empty(self, node_id: int) -> bool:
         """ Returns a bool of the contents of a cell. """
         return not self.G.nodes[node_id]["agent"]

tests/test_space.py

@@ -357,6 +357,15 @@ def test_move_agent(self):
         assert _agent not in self.space.G.nodes[initial_pos]["agent"]
         assert _agent in self.space.G.nodes[final_pos]["agent"]
 
+    def test_remove_agent(self):
+        for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE):
+            a = self.agents[i]
+            assert a.pos == pos
+            assert a in self.space.G.nodes[pos]["agent"]
+            self.space.remove_agent(a)
+            assert a.pos is None
+            assert a not in self.space.G.nodes[pos]["agent"]
+
     def test_is_cell_empty(self):
         assert not self.space.is_cell_empty(0)
         assert self.space.is_cell_empty(TestSingleNetworkGrid.GRAPH_SIZE - 1)