diff --git a/script/pgn_to_plain.py b/script/pgn_to_plain.py index 5f9300cb..596df5f3 100644 --- a/script/pgn_to_plain.py +++ b/script/pgn_to_plain.py @@ -1,11 +1,15 @@ import chess.pgn import argparse import glob +import re from typing import List # todo close in c++ tools using pgn-extract # https://www.cs.kent.ac.uk/people/staff/djb/pgn-extract/help.html#-w +commentRe = re.compile("([+-]*M*[0-9.]*)/([0-9]*)") +mateRe = re.compile("([+-])M([0-9]*)") + def parse_result(result_str:str, board:chess.Board) -> int: if result_str == "1/2-1/2": return 0 @@ -20,7 +24,7 @@ def parse_result(result_str:str, board:chess.Board) -> int: else: return -1 else: - print("illeagal result", result_str) + print("illegal result", result_str) raise ValueError def game_sanity_check(game: chess.pgn.Game) -> bool: @@ -28,21 +32,52 @@ def game_sanity_check(game: chess.pgn.Game) -> bool: print("invalid result", game.headers["Result"]) return False return True + +def parse_comment_for_score(comment_str: str, board: chess.Board) -> int: + global commentRe + global mateRe + + try: + m = commentRe.search(comment_str) + if m: + score = m.group(1) + # depth = int(m.group(2)) + m = mateRe.search(score) + if m: + if m.group(1) == "+": + score = 32000 - int(m.group(2)) + else: + score = -32000 + int(m.group(2)) + else: + score = int(float(score) * 208) # pawn to SF PawnValueEg + + if board.turn == chess.BLACK: + score = -score + else: + score = 0 + except: + score = 0 + + return score def parse_game(game: chess.pgn.Game, writer, start_play: int=1)->None: board: chess.Board = game.board() if not game_sanity_check(game): return + result: str = game.headers["Result"] - for ply, move in enumerate(game.mainline_moves()): + ply = 0 + for node in game.mainline(): + move = node.move if ply >= start_play: + comment: str = node.comment writer.write("fen " + board.fen() + "\n") writer.write("move " + str(move) + "\n") - writer.write("score 0\n") + writer.write("score " + str(parse_comment_for_score(comment, board)) + "\n") writer.write("ply " + str(ply)+"\n") writer.write("result " + str(parse_result(result, board)) +"\n") writer.write("e\n") - + ply += 1 board.push(move) def main(): @@ -53,6 +88,7 @@ def main(): args = parser.parse_args() pgn_files: List[str] = glob.glob(args.pgn) + pgn_files = sorted(pgn_files, key=lambda x:float(re.findall("-(\d+).pgn",x)[0] if re.findall("-(\d+).pgn",x) else 0.0)) f = open(args.output, 'w') for pgn_file in pgn_files: print("parse", pgn_file)