My submission for the Thumbtack PyCon schwag

connect4.py

#!/usr/bin/env python
import json
import sys


# Set up directions.
# ==================
# a b c
# d   e
# f g h

class Direction(object):
    def __init__(self, name, dx, dy):
        self.name = name # for debugging
        self.dx = dx
        self.dy = dy

a = Direction('a', -1, -1)
b = Direction('b', 0, -1)
c = Direction('c', 1, -1)
d = Direction('d', -1, 0)
e = Direction('e', 1, 0)
f = Direction('f', -1, 1)
g = Direction('g', 0, 1)
h = Direction('h', 1, 1)
_directions = [a,b,c,d,e,f,g,h]

def directions(board, _y, _x):
    """Given starting coordinates, yield 9 lists.

    Each list represents the sequence of contiguous pieces for the player with
    the piece at _y, _x. We'll check the length in the caller to determine the
    winner. Since this is a generator, we don't compute more sequences than we
    have to in order to determine the winner.

    """
    target = board[_y][_x]
    for direction in _directions:
        y = _y
        x = _x
        pieces = [target]
        while 1:
            y += direction.dy 
            x += direction.dx
            try:
                next = board[y][x]
            except IndexError:
                break
            if next != target:
                break
            pieces.append(next)
        yield pieces 


def play(board):
    """Given a 2-d list, return a string from {".", "O", "X"}.

    We discover the winner via depth-first search.

    """
    for y in range(len(board)):
        for x in range(len(board[y])):
            player = board[y][x]
            if player in ["O", "X"]:
                for pieces in directions(board,y,x):
                    if len(pieces) == 4:
                        return player # Our work here is done.
    return "."


# Tests.
# ======

game1 = """
[
  [".", ".", ".", ".", ".", ".", "."],
  [".", ".", ".", ".", ".", ".", "."],
  [".", ".", "O", ".", ".", ".", "."],
  [".", ".", "X", "O", "X", "X", "."],
  [".", ".", "X", "X", "X", "O", "X"],
  [".", ".", "O", "X", "X", "O", "X"]
]"""

game2 = """
[
  [".", ".", ".", ".", ".", ".", "."],
  [".", ".", ".", ".", ".", ".", "."],
  [".", ".", "O", ".", ".", ".", "."],
  [".", ".", "X", "O", "X", "X", "."],
  [".", ".", "X", "X", "O", "O", "X"],
  [".", ".", "O", "X", "X", "O", "X"]
]
"""

game3 = """
[
  [".", ".", ".", ".", ".", ".", "."],
  [".", ".", ".", ".", ".", ".", "."],
  [".", "O", "O", ".", ".", ".", "."],
  [".", "O", "X", "O", "X", "X", "."],
  [".", "O", "X", "X", "X", "O", "X"],
  [".", "O", "O", "X", "X", "O", "X"]
]
"""

game4 = """
[
  [".", ".", ".", ".", ".", ".", "."],
  [".", ".", ".", ".", ".", ".", "."],
  [".", ".", ".", "X", ".", ".", "."],
  [".", ".", "X", "O", "X", "X", "."],
  [".", "X", "X", "X", "O", "O", "X"],
  ["X", "X", "O", "X", "X", "O", "X"]
]
"""

games = ((game1, "."), (game2, "O"), (game3, "O"), (game4, "X"))

def test():
    for i, (game, expected) in enumerate(games):
        actual = play(json.loads(game))
        assert actual == expected, ("%s should have won game %d, not %s." 
                                    % (expected, i+1, actual))
        sys.stdout.write('.'); sys.stdout.flush()
    print


# main
# ====

def main(argv):
    if "--test" in argv:
        test()
    else:
        game = json.load(sys.stdin)
        result = play(game)
        if result in ["O", "X"]:
            print "Winner:", result
        else:
            assert result == "." # sanity check
            print "No winner."

if __name__ == "__main__":
    main(sys.argv)