Repository: projectmesa/mesa
Pull request #820 (2020-04-25 20:16:38)

Advanced indexing of grid
Currently mesa only allows direct grid access via `grid[x][y]`

This PR advances the possibilities to the following forms:

1. `grid[x][y]`
1. `grid[x, y]`
2. `grid[:, :]`
4. `grid[(x1, y1), (x2, y2), ...]`

where (1) is the current way, (2) is a tuple of coordinates (the usual `pos` attribute of agents) and (3) is a sliced version of (2): It allows to quickly select a single column (`grid[1, :]`), row (`grid[:, 2]`) or any other rectangular shape/neighborhood (`grid[1:3, 2:7:2]`). 

(4) accepts an sequence of positions, so for example something like that would work:
```python
neighborhood = grid.get_neighborhood(pos)
content = grid[neighborhood]
```
The implementation differs from `get_cell_list_contents` in that it will always return a list of the same length as the input, whereas `get_cell_list_contents` only returns the content of non-empty cells. 

This PR is not associated with any issue, but aims to ease grid access. 


 

Corvince: Travis currently fails, because of the @overload decorator used that makes flake8 cry. However it actually cries because of pyflakes, which is pinned to a version `<2.2.0` via flake8. This should be updated in the next version of flake8 (the [PR for that issue](https://gitlab.com/pycqa/flake8/-/merge_requests/417) is already merged) 

Patch:
diff --git a/mesa/space.py b/mesa/space.py
index 1b0957de..43078574 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -17,8 +17,21 @@ import itertools
 
 import numpy as np
 
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union
-from mesa.agent import Agent
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+    overload,
+)
+from .agent import Agent
 
 Coordinate = Tuple[int, int]
 GridContent = Union[Optional[Agent], Set[Agent]]
@@ -105,8 +118,69 @@ class Grid:
         """ Default value for new cell elements. """
         return None
 
+    @overload
     def __getitem__(self, index: int) -> List[GridContent]:
-        return self.grid[index]
+        ...
+
+    @overload
+    def __getitem__(
+        self, index: Tuple[Union[int, slice], Union[int, slice]]
+    ) -> Union[GridContent, List[GridContent]]:
+        ...
+
+    @overload
+    def __getitem__(self, index: Sequence[Coordinate]) -> List[GridContent]:
+        ...
+
+    def __getitem__(
+        self,
+        index: Union[
+            int, Sequence[Coordinate], Tuple[Union[int, slice], Union[int, slice]],
+        ],
+    ) -> Union[GridContent, List[GridContent]]:
+        """Access contents from the grid."""
+
+        if isinstance(index, int):
+            # grid[x]
+            return self.grid[index]
+
+        if isinstance(index[0], tuple):
+            # grid[(x1, y1), (x2, y2)]
+            index = cast(Sequence[Coordinate], index)
+
+            cells = []
+            for pos in index:
+                x1, y1 = self.torus_adj(pos)
+                cells.append(self.grid[x1][y1])
+            return cells
+
+        x, y = index
+
+        if isinstance(x, int) and isinstance(y, int):
+            # grid[x, y]
+            index = cast(Coordinate, index)
+            x, y = self.torus_adj(index)
+            return self.grid[x][y]
+
+        if isinstance(x, int):
+            # grid[x, :]
+            x, _ = self.torus_adj((x, 0))
+            x = slice(x, x + 1)
+
+        if isinstance(y, int):
+            # grid[:, y]
+            _, y = self.torus_adj((0, y))
+            y = slice(y, y + 1)
+
+        # grid[:, :]
+        x, y = (cast(slice, x), cast(slice, y))
+        cells = []
+        for rows in self.grid[x]:
+            for cell in rows[y]:
+                cells.append(cell)
+        return cells
+
+        raise IndexError
 
     def __iter__(self) -> Iterator[GridContent]:
         """


Test patch:
diff --git a/tests/test_grid.py b/tests/test_grid.py
index f7be0f6f..1079fc06 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -431,5 +431,31 @@ class TestHexGridTorus(TestBaseGrid):
         assert len(neighborhood) == 6
 
 
-if __name__ == "__main__":
+class TestIndexing:
+    # Create a grid where the content of each coordinate is a tuple of its coordinates
+    grid = Grid(3, 5, True)
+    for _, x, y in grid.coord_iter():
+        grid.grid[x][y] = (x, y)
+
+    def test_int(self):
+        assert self.grid[0][0] == (0, 0)
+
+    def test_tuple(self):
+        assert self.grid[1, 1] == (1, 1)
+
+    def test_list(self):
+        assert self.grid[(0, 0), (1, 1)] == [(0, 0), (1, 1)]
+        assert self.grid[(0, 0), (5, 3)] == [(0, 0), (2, 3)]
+
+    def test_torus(self):
+        assert self.grid[3, 5] == (0, 0)
+
+    def test_slice(self):
+        assert self.grid[:, 0] == [(0, 0), (1, 0), (2, 0)]
+        assert self.grid[::-1, 0] == [(2, 0), (1, 0), (0, 0)]
+        assert self.grid[1, :] == [(1, 0), (1, 1), (1, 2), (1, 3), (1, 4)]
+        assert self.grid[:, :] == [(x, y) for x in range(3) for y in range(5)]
+
+
+if __name__ == '__main__':
     unittest.main()