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, }