From 384a7cee6d06206ab27cfbb828b04cc33d26be67 Mon Sep 17 00:00:00 2001
From: corvince <13568919+Corvince@users.noreply.github.com>
Date: Sun, 15 Sep 2024 22:55:04 +0200
Subject: [PATCH] update dict-key handling

---
 mesa/experimental/cell_space/cell.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mesa/experimental/cell_space/cell.py b/mesa/experimental/cell_space/cell.py
index c06afe3d107..97c34abb48a 100644
--- a/mesa/experimental/cell_space/cell.py
+++ b/mesa/experimental/cell_space/cell.py
@@ -74,7 +74,7 @@ def connect(self, other: Cell, name: str | None = None) -> None:
         """
         if name is None:
             name = str(other.coordinate)
-        self.connections.update({name: other})
+        self.connections[name] = other
 
     def disconnect(self, other: Cell) -> None:
         """Disconnects this cell from another cell.
@@ -83,7 +83,9 @@ def disconnect(self, other: Cell) -> None:
             other (Cell): other cell to remove from connections
 
         """
-        self.connections = {k: v for k, v in self.connections.items() if v != other}
+        keys_to_remove = [k for k, v in self.connections.items() if v == other]
+        for key in keys_to_remove:
+            del self.connections[key]
 
     def add_agent(self, agent: CellAgent) -> None:
         """Adds an agent to the cell.