Skip to content

Commit

Permalink
restructure and rename
Browse files Browse the repository at this point in the history
  • Loading branch information
Corvince committed Sep 12, 2024
1 parent ff611c9 commit ec36c08
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 16 deletions.
6 changes: 3 additions & 3 deletions mesa/experimental/cell_space/cell.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down
52 changes: 39 additions & 13 deletions mesa/experimental/cell_space/cell_agent.py
Original file line number Diff line number Diff line change
@@ -1,24 +1,52 @@
from __future__ import annotations

from typing import TYPE_CHECKING, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Any, Protocol, TypeVar

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
Expand All @@ -34,10 +62,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)
Expand Down

0 comments on commit ec36c08

Please sign in to comment.