Simple game of life implementation - Python3 + pygame

gistfile1.txt

import pygame

W, H = 32, 24

width, height = 800, 600

def init():
    screen = pygame.display.set_mode((width, height))
    return screen


def make_data(size):
    W, H = size
    data = []
    for y in range(H):
        data.append([0 for x in range(W)])
    return data



class Grid:

    color = (255, 255, 255)

    def __init__(self):
        self.data = make_data((W, H))

    def __getitem__(self, item):
        try:
            return self.data[item[0]][item[1]]
        except IndexError:
            return 0

    def __setitem__(self, item, value):
        self.data[item[0]][item[1]] = value

    def count_neighbours(self, pos):
        x, y = pos
        count = sum([
            self[x - 1, y - 1], self[x, y - 1], self[x + 1, y - 1],
            self[x - 1, y], self[x + 1, y],
            self[x - 1, y + 1], self[x, y + 1], self[x + 1, y + 1]
        ])
        return count

    def update(self):
        new_data = make_data((W,H))
        for y in range(H):
            for x in range(W):
                n = self.count_neighbours((y, x))
                if n == 2 and self[y,x]:
                    new_data[y][x] = 1
                elif n == 3:
                    new_data[y][x] = 1
                else:
                    new_data[y][x] = 0
        self.data = new_data

    def draw(self, screen):
        cs = cell_size = width / W

        for i, row in enumerate(self.data):
            for j, cell in enumerate(row):
                pygame.draw.rect(screen, (0,0,0), (j * cs, i * cs, cs, cs))
                if cell:
                    pygame.draw.rect(screen, self.color, (j * cs + 1 , i * cs + 1 , cs - 2, cs - 2))


def setup(grid):
    grid[10, 10] = 1
    grid[10, 11] = 1
    grid[10, 12] = 1
    grid[11, 12] = 1


def main():
    try:
        screen = init()
        grid = Grid()
        setup(grid)
        while True:
            pygame.event.pump()
            keys = pygame.key.get_pressed()
            if keys[pygame.K_ESCAPE]:
                break

            grid.update()
            grid.draw(screen)
            pygame.display.flip()
            pygame.time.delay(200)

    finally:
        pygame.quit()

if __name__ == "__main__":
    main()