connect4.py - play connect 4 in the terminal.

connect4.py

WIDTH = 7
HEIGHT = 6
NULL = " "

def connected_4(l):
    # Tests whether or not the list, l, has 
    # four similar characters in a row.
    last = None
    run = 1
    for item in l:
        if item == NULL:
            last = item
            run = 1
        elif item == last:
            run += 1
        else:
            last = item
            run = 1
        if run >= 4:
            return item
    return False

def sheer(board, direction):
    """
    Sheers the board so that...
            ....
            ....
            ....
    becomes
            .             .
            ..           ..
            ...         ...
             ...   or  ...
              ..       ..
               .       .

    and yields each row.

    direction should be +1 or -1
    """
    for ix in range(0, WIDTH+HEIGHT-1):
        row = []
        rowx = min(ix, HEIGHT-1)    # 0, 1, 2, 2, 2, ...
        if direction == 1:
            colx = ix - rowx          # 0, 0, 0, 1, 2, ...
        else:
            colx = (WIDTH-1) - (ix - rowx)
        while 0 <= rowx < HEIGHT and 0 <= colx < WIDTH:
            row.append(board[rowx][colx])
            rowx -= 1
            colx += (1 * direction)
        yield row


def test_sheer():
    # Just to test my code works. 
    board = [ [chr(65+x) for x in range(WIDTH)] for y in range(HEIGHT) ]
    board[0][0] = "X"
    
    import pprint
    pprint.pprint(board)
    print "\n\n"
    for k in sheer(board, -1):
        print k
    print "\n\n"
    for k in sheer(board, -1):
        print k


class GameBoard(object):

    def __init__(self):
        self.board = [ [" " for x in range(WIDTH)] for y in range(HEIGHT) ]
        self.col_fullness = [0 for x in range(WIDTH)]

    def printboard(self):
        print "--"*WIDTH
        print " ".join(str(n) for n in range(WIDTH))
        for row in reversed(self.board):
            print " ".join(row)
        print "--"*WIDTH

    def winner(self):
        for row in self.board:
            winner = connected_4(row)
            if winner:
                return winner

        for col in range(WIDTH):
            winner = connected_4( [row[col] for row in self.board] )
            if winner:
                return winner

        # Test diagonals.
        for row in sheer(self.board,1):
            winner = connected_4(row)
            if winner:
                return winner

        for row in sheer(self.board,-1):
            winner = connected_4(row)
            if winner:
                return winner

        return False

    def fullboard(self):
        return all((self.col_fullness[ix] == HEIGHT) for ix in range(WIDTH))

    def play(self, token, column):
        try:
            column = int(column)
        except:
            print "Enter an integer"
            return False
        if column < 0 or column >= WIDTH:
            return False
        if self.col_fullness[column] >= HEIGHT:
            return False

        self.board[self.col_fullness[column]][column] = token
        self.col_fullness[column] += 1
        return True


if __name__ == '__main__':
    gb = GameBoard()
    player = 'x'
    while (not gb.winner() and not gb.fullboard()):
        gb.printboard()
        play_success = False
        while not play_success:
            col_play = raw_input("Player '%s' take your turn:" % player)
            play_success = gb.play(player,col_play)

        if player == 'x':
            player = 'o'
        else:
            player = 'x'

    gb.printboard()
    winner = gb.winner()
    if winner:
        print "Yay! %s won!" % winner
    else:
        print "Stalemate!"