restructure and rename

Author: Corvince

mesa/experimental/cell_space/cell.py

@@ -7,7 +7,7 @@
 from mesa.experimental.cell_space.cell_collection import CellCollection
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space.cell_agent import CellHolder
+    from mesa.experimental.cell_space.cell_agent import DiscreteSpaceAgent
 
 
 class Cell:
@@ -80,7 +80,7 @@ def disconnect(self, other: Cell) -> None:
         """
         self._connections.remove(other)
 
-    def add_agent(self, agent: CellHolder) -> None:
+    def add_agent(self, agent: DiscreteSpaceAgent[Cell]) -> None:
         """Adds an agent to the cell.
 
         Args:
@@ -96,7 +96,7 @@ def add_agent(self, agent: CellHolder) -> None:
 
         self.agents.append(agent)
 
-    def remove_agent(self, agent: CellHolder) -> None:
+    def remove_agent(self, agent: DiscreteSpaceAgent[Cell]) -> None:
         """Removes an agent from the cell.
 
         Args:

mesa/experimental/cell_space/cell_agent.py

@@ -1,24 +1,53 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
 
+from mesa.agent import Agent
 from mesa.experimental.cell_space.discrete_space import DiscreteSpace
 
 if TYPE_CHECKING:
-    from mesa.experimental.cell_space import Cell, Grid
+    from mesa.experimental.cell_space import Cell
 
+T = TypeVar("T", bound="Cell")
 
-@runtime_checkable
-class CellHolder(Protocol):
-    cell: Cell | None
+
+class DiscreteSpaceAgent(Protocol[T]):
+    cell: T | None
+    space: DiscreteSpace[T]
+
+    def move_to(self, cell: T) -> None:
+        ...
+
+    def move_relative(self, directions: tuple[int, ...], distance: int = 1):
+        ...
 
 
 class CellAgent:
-    cell: Cell | None
-    space: DiscreteSpace[Cell]
-    """Cell Agent is an extension of the Agent class and adds behavior for moving in discrete spaces"""
+    """Cell Agent is an Agent class that adds behavior for moving in discrete spaces
+
+    Attributes:
+        space (DiscreteSpace): the discrete space the agent is in
+        cell (Cell): the cell the agent is in
+    """
 
-    def move_to(self: CellHolder, cell: Cell) -> None:
+    def __init__(
+        self,
+        space: DiscreteSpace[Cell],
+        cell: Cell | None = None,
+        *args: tuple[Any],
+        **kwargs: dict[str, Any],
+    ):
+        super().__init__(*args, **kwargs)
+        self.space = space
+        self.cell = cell
+        if cell is not None:
+            cell.add_agent(self)
+
+    @property
+    def coordinate(self) -> tuple[int, ...]:
+        return self.cell.coordinate if self.cell else ()
+
+    def move_to(self, cell: Cell) -> None:
         if self.cell is not None:
             self.cell.remove_agent(self)
         self.cell = cell
@@ -34,10 +63,8 @@ def move_relative(self, directions: tuple[int, ...], distance: int = 1):
         self.move_to(new_cell)
 
 
-class Grid2DMovingAgent(CellAgent):
-    grid: Grid[Cell]
-
-    def move(self, direction: str, distance: int = 1):
+class Grid2DMovingAgent:
+    def move(self: DiscreteSpaceAgent[Cell], direction: str, distance: int = 1):
         match direction:
             case "N" | "North" | "Up":
                 self.move_relative((-1, 0), distance)