HomeBlogThe snake, seen from the eyes of the CamelChess::Rep in Python

Chess::Rep in Python

The Perl module can be found at CPAN.  I'm pasting the Python code below in case you want to take a quick look.

If you want to run the sample, you can download both here: chesstest.tar.bz2.  Unpack and run test.py or test.pl in the Python or Perl subdirectories.  They both parse a large (1.4MB) file containing lots of Chess games.  In my tests, the Python version takes around 5.5 minutes, while the Perl version takes like 11 minutes.  If you have any optimization ideas, please let me know.

Of course, feel free to use this code under the same terms as, err, Python itself. ;-)  It's pretty useful.  Note that it's not a PGN parser in itself.  See the documentation of the Perl module for details.

import re, string, struct

CASTLE_W_OO  = 1
CASTLE_W_OOO = 2
CASTLE_B_OO  = 4
CASTLE_B_OOO = 8

PIECE_TO_ID = {
    'p' : 0x01,              # black pawn
    'n' : 0x02,              # black knight
    'k' : 0x04,              # black king
    'b' : 0x08,              # black bishop
    'r' : 0x10,              # black rook
    'q' : 0x20,              # black queen
    'P' : 0x81,              # white pawn
    'N' : 0x82,              # white knight
    'K' : 0x84,              # white king
    'B' : 0x88,              # white bishop
    'R' : 0x90,              # white rook
    'Q' : 0xA0,              # white queen
}

ID_TO_PIECE = [
    0,                       # 0
    'p',                     # 1
    'n',                     # 2
    0,                       # 3
    'k',                     # 4
    0,                       # 5
    0,                       # 6
    0,                       # 7
    'b',                     # 8
    0,                       # 9
    0,                       # 10
    0,                       # 11
    0,                       # 12
    0,                       # 13
    0,                       # 14
    0,                       # 15
    'r',                     # 16
    0,                       # 17
    0,                       # 18
    0,                       # 19
    0,                       # 20
    0,                       # 21
    0,                       # 22
    0,                       # 23
    0,                       # 24
    0,                       # 25
    0,                       # 26
    0,                       # 27
    0,                       # 28
    0,                       # 29
    0,                       # 30
    0,                       # 31
    'q',                     # 32
]

FEN_STANDARD = 'rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1'

MOVES_N = 31, 33, 14, 18, -18, -14, -33, -31
MOVES_B = 15, 17, -15, -17
MOVES_R = 1, 16, -16, -1
MOVES_K = MOVES_B + MOVES_R

GO_MOVE_REGEXPS = {
    "piece" : re.compile(r"[PNBRQK]"),
    "full"  : re.compile(r"([a-h][1-8])[:x-]?([a-h][1-8])", re.I),
    "col"   : re.compile(r"([a-h])[:x-]?([a-h][1-8])", re.I),
    "row"   : re.compile(r"([1-8])[:x-]?([a-h][1-8])", re.I),
    "to"    : re.compile(r"[:x-]?([a-h][1-8])", re.I),
    "promo" : re.compile(r"=?([RNBQ])", re.I),
    }

# -----[ exceptions ]----- #

class InvalidMoveEx(Exception):
    def __init__(self, msg):
        self.message = msg
    def __str__(self):
        return repr(self.message)

# -----[ public functions ]----- #

def get_index_from_row_col(row, col):
    """ Returns an index from row, col """
    return row << 4 | col

def get_row_col(row, col = None):
    """ Returns row, col given an index or a field ID """
    if col == None:
        if isinstance(row, int):
            return (row & 0x70) >> 4, row & 0x07
        else:
            # row is an ID
            (col, row) = struct.unpack('bb', row.upper())
            return row - 49, col - 65
    return row, col

def get_field_id(row, col = None):
    (row, col) = get_row_col(row, col)
    return struct.pack('bb', col + 97, row + 49)

def get_piece_rep(piece):
    rep = ID_TO_PIECE[ piece & 0x3F ]
    if piece & 0x80:
        rep = rep.upper()
    return rep

def get_index(id):
    (row, col) = get_row_col(id)
    return get_index_from_row_col(row, col)

# -----[ public classes ]----- #

class Board:

    def __init__(self, fen = FEN_STANDARD):
        self.set_from_fen(fen)

    def _reset(self):
        self.pos = [ 0 ] * 128
        self.castle = CASTLE_W_OO | CASTLE_W_OOO | CASTLE_B_OO | CASTLE_B_OOO
        self.has_castled = 0
        self.to_move = 0x80
        self.enpa = None
        self.halfmove = 0
        self.fullmove = 0
        self.status = 0

    def reset(self):
        self.set_from_fen(FEN_STANDARD)

    def set_from_fen(self, fen):
        self._reset()
        (board, to_move, castle, enpa, halfmove, fullmove) = re.split(r'\s+', fen)
        board = re.split(r'/', board)
        board.reverse()
        for row, line in enumerate(board):
            col = 0
            for p in line:
                if p in PIECE_TO_ID:
                    self.set_piece_at_index(get_index_from_row_col(row, col), PIECE_TO_ID[p])
                    col += 1
                elif p in "12345678":
                    col += int(p)
        c = 0
        if 'K' in castle: c |= CASTLE_W_OO
        if 'Q' in castle: c |= CASTLE_W_OOO
        if 'k' in castle: c |= CASTLE_B_OO
        if 'q' in castle: c |= CASTLE_B_OOO
        self.castle = c
        to_move = to_move.lower()
        if to_move == 'w':
            self.to_move = 0x80
        elif to_move == 'b':
            self.to_move = 0
        else:
            self.to_move = None
        self.enpa = (get_index(enpa) if enpa != '-' else 0)
        self.fullmove = int(fullmove)
        self.halfmove = int(halfmove)
        self.compute_valid_moves()

    def set_piece_at_index(self, index, p):
        old = self.pos[index]
        self.pos[index] = p
        return old

    def get_piece_at_index(self, index):
        return self.pos[index]

    def get_fen(self, short = False):
        a = []
        for row in reversed(range(8)):
            line = ''
            empty = 0
            for col in range(8):
                p = self.get_piece_at_index(get_index_from_row_col(row, col))
                if p:
                    p = (ID_TO_PIECE[p & 0x3F].upper() if p & 0x80 else ID_TO_PIECE[p])
                    if empty:
                        line += str(empty)
                    empty = 0
                    line += p
                else:
                    empty += 1
            if empty:
                line += str(empty)
            a.append(line)
        a = [ '/'.join(a) ]
        castle = self.castle
        c = ''
        if castle & CASTLE_W_OO  : c += 'K'
        if castle & CASTLE_W_OOO : c += 'Q'
        if castle & CASTLE_B_OO  : c += 'k'
        if castle & CASTLE_B_OOO : c += 'q'
        a.extend([ 'w' if self.to_move else 'b',
                   c or '-',
                   get_field_id(self.enpa) if self.enpa else '-' ])
        if not short:
            a.extend([ str(self.halfmove), str(self.fullmove) ])
        return ' '.join(a)

    def can_castle(self, color, ooo):
        if color:
            return self.castle & (CASTLE_W_OOO if ooo else CASTLE_W_OO)
        else:
            return self.castle & (CASTLE_B_OOO if ooo else CASTLE_B_OO)

    def dump_pos(self):
        fen = self.get_fen().split(' ')[0]
        fen = re.sub(r'[1-8]', lambda x: ' ' * int(x.group()), fen)
        fen = re.sub(r'[^/]', lambda x: '|' + x.group(), fen)
        fen = re.sub(r'/', "|\n|-+-+-+-+-+-+-+-|\n", fen)
        fen += "|"
        return fen

    def dump_status(self):
        moves = self.status["moves"]
        print "* Possible moves:"
        for m in moves:
            piece = get_piece_rep(m["piece"])
            args = {
                'p'    : piece,
                'from' : get_field_id(m['from']),
                'to'   : ', '.join(map(get_field_id, m['to']))
                }
            print "[%(p)s] from %(from)s can move to: %(to)s" % args

        print "* Moves by target field and by piece:"
        for (key, val) in self.status["type_moves"].items():
            print get_field_id(key), " --- "
            for (piece, a) in val.items():
                for tmp in a:
                    print "  : ", get_piece_rep(piece), " from ", get_field_id(tmp)

    def is_attacked(self, i, op_color = None, try_move = None):
        if op_color == None:
            op_color = self.to_move ^ 0x80

        def test(type, i):
            if i & 0x88:
                return 1
            p = self.pos[i]
            # print "Here %(i)s %(p)02X looking for %(type)02X %(op)d" % { 'i': get_field_id(i), 'p': p, 'op': op_color, 'type': type }
            if try_move:
                if i == try_move['from']:
                    p = 0
                elif i == try_move['to']:
                    p = try_move['piece']
            if p and p & type and p & 0x80 == op_color:
                raise _IsAttacked()
            return p

        try:

            # check pawns
            if op_color:
                test(0x01, i - 15)
                test(0x01, i - 17)
            else:
                test(0x01, i + 15)
                test(0x01, i + 17)

            # check knights
            for step in MOVES_N:
                test(0x02, i + step)

            # check bishops or queens
            for step in MOVES_B:
                j = i + step
                while not test(0x28, j):
                    j += step

            # check rooks or queens
            for step in MOVES_R:
                j = i + step
                while not test(0x30, j):
                    j += step

            # check opponent king
            for step in MOVES_K:
                test(0x04, i + step)

        except _IsAttacked:
            return True

        return False

    def compute_valid_moves(self):
        pieces = []
        king = None
        op_color = self.to_move ^ 0x80

        for row in range(8):
            for col in range(8):
                i = get_index_from_row_col(row, col)
                p = self.get_piece_at_index(i)
                if p and p & 0x80 == self.to_move:
                    pieces.append({ 'from' : i, 'piece' : p })
                    if p & 0x04:
                        king = i

        if king:
            self.in_check = self.is_attacked(king, op_color)

        all_moves = []
        hash_moves = {}
        type_moves = {}

        for p in pieces:
            index = p['from']
            moves = _get_allowed_moves(self, index)
            piece = p['piece']
            if king:
                is_king = index == king
                try_move = { 'from'  : index,
                             'piece' : piece }
                def valid(x):
                    if x & 0x100:
                        return True
                    try_move['to'] = x
                    return not self.is_attacked(x if is_king else king, op_color, try_move)
                valid_moves = filter(valid, moves)
            else:
                valid_moves = moves
            if len(valid_moves) > 0:
                p['to'] = valid_moves
                all_moves.append(p)
                for to in valid_moves:
                    to = to & 0xFF
                    hash_moves[str(index) + "-" + str(to)] = 1
                    if to not in type_moves:
                        type_moves[to] = {}
                    a = type_moves[to]
                    if piece not in a:
                        a[piece] = []
                    a[piece].append(index)

        self.status = { "moves"      : all_moves,
                        "pieces"     : pieces,
                        "hash_moves" : hash_moves,
                        "type_moves" : type_moves,
                        "check"      : self.in_check,
                        "mate"       : self.in_check and not all_moves,
                        "stalemate"  : not self.in_check and not all_moves }

        return self.status

    def go_move(self, move):

        color = self.to_move
        piece = None
        from_, to_ = None, None
        from_idx, to_idx = None, None
        promote = None
        promote_id = None
        san = None
        orig_move = move
        row = None
        col = None

        # check for queenside castle
        if move.find("O-O-O") == 0:
            move = "e1c1" if color else "e8c8"
        elif move.find("O-O") == 0:
            move = "e1g1" if color else "e8g8"
        elif GO_MOVE_REGEXPS["piece"].match(move):
            piece = move[0:1].lower()
            move = move[1:]

        while True:
            m = GO_MOVE_REGEXPS["full"].match(move)
            if m:
                from_, to_ = m.groups()
                move = move[m.end():]
                break

            m = GO_MOVE_REGEXPS["col"].match(move)
            if m:
                col = struct.unpack("b", m.group(1).upper())[0] - 65 # OMG
                to_ = m.group(2).upper()
                move = move[m.end():]
                break

            m = GO_MOVE_REGEXPS["row"].match(move)
            if m:
                row = struct.unpack("b", m.group(1))[0] - 49
                to_ = m.group(2).upper()
                move = move[m.end():]
                break

            m = GO_MOVE_REGEXPS["to"].match(move)
            if m:
                to_ = m.group(1).upper()
                move = move[m.end():]
                break

            raise InvalidMoveEx("ERROR: couldn't parse move: %(move)s" % { "move": orig_move })

        m = GO_MOVE_REGEXPS["promo"].match(move)
        if m:
            promote = m.group(1).upper()

        if piece:
            piece = PIECE_TO_ID[piece]
        else:
            if not from_:
                piece = 1               # black pawn
            else:
                from_idx = get_index(from_)
                piece = self.get_piece_at_index(from_idx)
                if not piece:
                    die("ERROR: illegal move %(move)s (field %(from)s is empty)" % { "move": move, "from": from_ })

        piece = piece | color

        if not to_:
            raise InvalidMoveEx("Can't parse move %(move)s (missing target field)" % { "move": orig_move })

        to_idx = get_index(to_)

        try:
            tpmove = self.status["type_moves"][to_idx][piece]
        except KeyError:
            raise InvalidMoveEx("Illegal move %(move)s - no piece can move to %(to)s" % { "move": orig_move, "to": to_ })

        if not tpmove:
            raise InvalidMoveEx("Illegal move %(move)s" % { "move": orig_move })

        if not from_:
            if len(tpmove) == 1:
                from_idx = tpmove[0]
            else:
                for origin in tpmove:
                    t_row, t_col = get_row_col(origin)
                    if row == t_row:
                        from_idx = origin
                        break
                    if col == t_col:
                        from_idx = origin
                        break

            if from_idx != None:
                from_ = get_field_id(from_idx)
            else:
                raise InvalidMoveEx("Ambiguous move %(move)s" % { "move": orig_move })

        elif from_idx not in tpmove:
            raise InvalidMoveEx("Illegal move %(move)s" % { "move": orig_move })

        if from_idx is None:
            from_idx = get_index(from_)

        from_row, from_col = get_row_col(from_idx)
        to_row, to_col = get_row_col(to_idx)

        prev_enpa = self.enpa
        self.enpa = 0
        is_capture = False
        is_pawn = piece & 0x01

        while True:
            # 1. if it's castling, we have to move the rook
            if piece & 0x04:            # king?
                if from_idx == 0x04 and to_idx == 0x06:
                    san = "O-O"
                    self.has_castled = self.has_castled | CASTLE_W_OO
                    self._move_piece(0x07, 0x05)
                    break
                if from_idx == 0x74 and to_idx == 0x76:
                    san = "O-O"
                    self.has_castled = self.has_castled | CASTLE_B_OO
                    self._move_piece(0x77, 0x75)
                    break
                if from_idx == 0x04 and to_idx == 0x02:
                    san = "O-O-O"
                    self.has_castled = self.has_castled | CASTLE_W_OOO
                    self._move_piece(0x00, 0x03)
                    break
                if from_idx == 0x74 and to_idx == 0x72:
                    san = "O-O-O"
                    self.has_castled = self.has_castled | CASTLE_B_OOO
                    self._move_piece(0x70, 0x73)
                    break

            # 2. is it en-passant?
            if is_pawn:
                if from_col != to_col and prev_enpa and prev_enpa == to_idx:
                    self.set_piece_at_index(get_index_from_row_col(from_row, to_col), 0)
                    is_capture = True
                    break
                if abs(from_row - to_row) == 2:
                    self.enpa = get_index_from_row_col((from_row + to_row) / 2, from_col)

            break

        if promote:
            promote_id = PIECE_TO_ID[ promote.lower() ] | color

        if self._move_piece(from_idx, to_idx, promote_id):
            is_capture = True

        self.to_move = self.to_move ^ 0x80

        if self.to_move:
            self.fullmove += 1

        if not is_pawn and not is_capture:
            self.halfmove += 1
        else:
            self.halfmove = 0

        status = self.compute_valid_moves()

        if not san:
            san = "" if is_pawn else ID_TO_PIECE[ piece & 0x3F ].upper()
            cnt = 1 if (is_capture and is_pawn or len(tpmove) > 1) else 0
            for origin in tpmove:
                if origin != from_idx and ((origin & 0x07) == (from_idx & 0x07)):
                    cnt = 2
                    break

            san += from_[:cnt]
            if is_capture:
                san += "x"
            san += to_.lower()
            if promote:
                san += "=" + promote

        if status["mate"]:
            san += "#"
        elif status["check"]:
            san += "+"

        return { "from"       : from_.lower(),
                 "from_index" : from_idx,
                 "from_row"   : from_row,
                 "from_col"   : from_col,
                 "to"         : to_.lower(),
                 "to_index"   : to_idx,
                 "to_row"     : to_row,
                 "to_col"     : to_col,
                 "piece"      : piece,
                 "promote"    : promote,
                 "san"        : san }

    def _move_piece(self, from_, to_, promote = None):
        p = self.set_piece_at_index(from_, 0)
        if p & 0x04:                    # is it king?
            if p & 0x08:
                self.castle = self.castle | CASTLE_W_OOO ^ CASTLE_W_OOO
                self.castle = self.castle | CASTLE_W_OO ^ CASTLE_W_OO
            else:
                self.castle = self.castle | CASTLE_B_OOO ^ CASTLE_B_OOO
                self.castle = self.castle | CASTLE_B_OO ^ CASTLE_B_OO
        elif from_ == 0x00 or to_ == 0x00:
            self.castle = self.castle | CASTLE_W_OOO ^ CASTLE_W_OOO
        elif from_ == 0x70 or to_ == 0x70:
            self.castle = self.castle | CASTLE_B_OOO ^ CASTLE_B_OOO
        elif from_ == 0x07 or to_ == 0x07:
            self.castle = self.castle | CASTLE_W_OO ^ CASTLE_W_OO
        elif from_ == 0x77 or to_ == 0x77:
            self.castle = self.castle | CASTLE_B_OO ^ CASTLE_B_OO
        return self.set_piece_at_index(to_, promote or p)

# -----[ private classes ]----- #

class _IsAttacked(Exception):
    pass

# -----[ private stuff ]----- #

def _add_if_valid(self, moves, _from, _to):

    if _to & 0x88:
        return None

    what = self.get_piece_at_index(_to)
    p = self.get_piece_at_index(_from)
    color = p & 0x80

    if p & 0x04 and self.is_attacked(_to):
        return None

    if not what:
        if p & 0x01:
            diff = (_from & 0x07) - (_to & 0x07)
            if diff < 0: diff = -diff
            if diff == 1:
                if self.enpa and _to == self.enpa:
                    moves.append(_to | 0x100)
                    return _to
                return None         # must take to move this way
        moves.append(_to)
        return _to

    if what & 0x80 != color:
        if p & 0x01 and _from & 0x07 == _to & 0x07:
            return None         # pawns can't take this way
        moves.append(_to)
        return _to

    return None

def _get_allowed_moves(self, index):
    p = ID_TO_PIECE[ self.get_piece_at_index(index) & 0x3F ].lower()
    method = GET_ALLOWED_MOVES[p]
    return method(self, index)

def _get_allowed_p_moves(self, index, moves = None):
    if moves == None: moves = []
    color = self.get_piece_at_index(index) & 0x80
    step = (16 if color else -16)
    not_moved = (index & 0xF0) == (0x10 if color else 0x60)
    if _add_if_valid(self, moves, index, index + step) != None and not_moved:
        _add_if_valid(self, moves, index, index + 2 * step)
    _add_if_valid(self, moves, index, index + (17 if color else -15))
    _add_if_valid(self, moves, index, index + (15 if color else -17))
    return moves

def _get_allowed_n_moves(self, index, moves = None):
    if moves == None: moves = []
    for step in MOVES_N:
        _add_if_valid(self, moves, index, index + step)
    return moves

def _get_allowed_b_moves(self, index, moves = None):
    if moves == None: moves = []
    for step in MOVES_B:
        i = index + step
        while _add_if_valid(self, moves, index, i):
            if self.get_piece_at_index(i):
                break
            i += step
    return moves

def _get_allowed_r_moves(self, index, moves = None):
    if moves == None: moves = []
    for step in MOVES_R:
        i = index + step
        while _add_if_valid(self, moves, index, i):
            if self.get_piece_at_index(i):
                break
            i += step
    return moves

def _get_allowed_q_moves(self, index, moves = None):
    if moves == None: moves = []
    _get_allowed_b_moves(self, index, moves)
    _get_allowed_r_moves(self, index, moves)
    return moves

def _get_allowed_k_moves(self, index, moves = None):
    if moves == None: moves = []
    color = self.get_piece_at_index(index) & 0x80
    for step in MOVES_K:
        if _add_if_valid(self, moves, index, index + step):
            if ( step == 1 and
                 not self.in_check and
                 self.can_castle(color, False) and
                 not self.get_piece_at_index(index + 1) and
                 not self.get_piece_at_index(index + 2) ):
                # kingside castling possible
                _add_if_valid(self, moves, index, index + 2)
            elif ( step == -1 and
                   not self.in_check and
                   self.can_castle(color, True) and
                   not self.get_piece_at_index(index - 1) and
                   not self.get_piece_at_index(index - 2) and
                   not self.get_piece_at_index(index - 3) ):
                # queenside castling
                _add_if_valid(self, moves, index, index - 2)
    return moves

GET_ALLOWED_MOVES = {
    'p': _get_allowed_p_moves,
    'n': _get_allowed_n_moves,
    'b': _get_allowed_b_moves,
    'r': _get_allowed_r_moves,
    'q': _get_allowed_q_moves,
    'k': _get_allowed_k_moves,
    }