code-retreat session 4

session4.py

from collections import Counter

class LiveCells(object):

    def __init__(self):
        self.cells = set()

    def add_cell(self, position):
        self.cells.add(position)
    
    def get_neighbours_count(self):
        result = Counter()
        for x, y in self.cells:
            for dx in range(-1, 2):
                for dy in range(-1, 2):
                    if dx == 0 and dy == 0:
                        continue
                    result[(x+dx, y+dy)] += 1
        for x, y in self.cells:
            if (x, y) in result:
                del result[(x, y)]
        return result

def test_live_cells_init():
    live_cells = LiveCells()
    assert len(live_cells.cells) == 0

    position = (0, 0)
    live_cells.add_cell(position)
    assert position in live_cells.cells

    assert live_cells.get_neighbours_count() == {
        (-1, -1): 1, (-1, 0): 1, (-1, 1): 1,
        (0, -1): 1, (0, 1): 1,
        (1, -1): 1, (1, 0): 1, (1, 1): 1}

    live_cells.add_cell((1, 0))
    neighbour_count = live_cells.get_neighbours_count()
    assert neighbour_count == {
        (-1, -1): 1, (-1, 0): 2, (-1, 1): 2, (-1, 2): 1,
        (0, -1):  1,                         (0, 2):  1,
        (1, -1):  1,  (1, 0): 2, (1, 1):  2, (1, 2):  1}