Repository: projectmesa/mesa

Issue:
baby patch/tweak: allow cell methods to take single cell
Cf. discussion in #176 -- taking on the small "(TODO: someone should probably fix that...)" from the tutorial re: grid's `get_cell_list_contents` requiring a list of cells even if it's just being passed a single cell.  It's a very easy fix, I'm on the case.  :-) 

Discussion:

Patch:
diff --git a/mesa/space.py b/mesa/space.py
index 5e3a9544..77fe8174 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -24,6 +24,20 @@ X = 0
 Y = 1
 
 
+def accept_tuple_argument(wrapped_function):
+    '''
+    Decorator to allow grid methods that take a list of (x, y) position tuples
+    to also handle a single position, by automatically wrapping tuple in
+    single-item list rather than forcing user to do it.
+    '''
+    def wrapper(*args):
+        if isinstance(args[1], tuple) and len(args[1]) == 2:
+            return wrapped_function(args[0], [args[1]])
+        else:
+            return wrapped_function(*args)
+    return wrapper
+
+
 class Grid(object):
     '''
     Base class for a square grid.
@@ -238,10 +252,11 @@ class Grid(object):
         x, y = pos
         return x < 0 or x >= self.width or y < 0 or y >= self.height
 
+    @accept_tuple_argument
     def iter_cell_list_contents(self, cell_list):
         '''
         Args:
-            cell_list: Array-like of (x, y) tuples
+            cell_list: Array-like of (x, y) tuples, or single tuple.
 
         Returns:
             A iterator of the contents of the cells identified in cell_list
@@ -249,10 +264,11 @@ class Grid(object):
         return (
             self[y][x] for x, y in cell_list if not self.is_cell_empty((x, y)))
 
+    @accept_tuple_argument
     def get_cell_list_contents(self, cell_list):
         '''
         Args:
-            cell_list: Array-like of (x, y) tuples
+            cell_list: Array-like of (x, y) tuples, or single tuple.
 
         Returns:
             A list of the contents of the cells identified in cell_list
@@ -418,10 +434,11 @@ class MultiGrid(Grid):
         x, y = pos
         self.grid[y][x].remove(agent)
 
+    @accept_tuple_argument
     def iter_cell_list_contents(self, cell_list):
         '''
         Args:
-            cell_list: Array-like of (x, y) tuples
+            cell_list: Array-like of (x, y) tuples, or single tuple.
 
         Returns:
             A iterator of the contents of the cells identified in cell_list


Test patch:
diff --git a/tests/test_grid.py b/tests/test_grid.py
index b558f4d9..c09f0496 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -53,6 +53,43 @@ class TestBaseGrid(unittest.TestCase):
             x, y = agent.pos
             assert self.grid[y][x] == agent
 
+    def test_cell_agent_reporting(self):
+        '''
+        Ensure that if an agent is in a cell, get_cell_list_contents accurately
+        reports that fact.
+        '''
+        for agent in self.agents:
+            x, y = agent.pos
+            assert agent in self.grid.get_cell_list_contents([(x, y)])
+
+    def test_listfree_cell_agent_reporting(self):
+        '''
+        Ensure that if an agent is in a cell, get_cell_list_contents accurately
+        reports that fact, even when single position is not wrapped in a list.
+        '''
+        for agent in self.agents:
+            x, y = agent.pos
+            assert agent in self.grid.get_cell_list_contents((x, y))
+
+    def test_iter_cell_agent_reporting(self):
+        '''
+        Ensure that if an agent is in a cell, iter_cell_list_contents
+        accurately reports that fact.
+        '''
+        for agent in self.agents:
+            x, y = agent.pos
+            assert agent in self.grid.iter_cell_list_contents([(x, y)])
+
+    def test_listfree_iter_cell_agent_reporting(self):
+        '''
+        Ensure that if an agent is in a cell, iter_cell_list_contents
+        accurately reports that fact, even when single position is not
+        wrapped in a list.
+        '''
+        for agent in self.agents:
+            x, y = agent.pos
+            assert agent in self.grid.iter_cell_list_contents((x, y))
+
     def test_neighbors(self):
         '''
         Test the base neighborhood methods on the non-toroid.