Enhance pgn_to_plain.py

in case a score can be parsed from the comment field in the pgn, add it to the output.
This form works for the fishtest pgns, and is quite common (cutechess-cli among others).
This commit is contained in:
Joost VandeVondele
2020-12-08 23:57:54 +01:00
committed by nodchip
parent d99ba07b81
commit 9c65e868f9

View File

@@ -1,11 +1,15 @@
import chess.pgn import chess.pgn
import argparse import argparse
import glob import glob
import re
from typing import List from typing import List
# todo close in c++ tools using pgn-extract # todo close in c++ tools using pgn-extract
# https://www.cs.kent.ac.uk/people/staff/djb/pgn-extract/help.html#-w # 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: def parse_result(result_str:str, board:chess.Board) -> int:
if result_str == "1/2-1/2": if result_str == "1/2-1/2":
return 0 return 0
@@ -20,7 +24,7 @@ def parse_result(result_str:str, board:chess.Board) -> int:
else: else:
return -1 return -1
else: else:
print("illeagal result", result_str) print("illegal result", result_str)
raise ValueError raise ValueError
def game_sanity_check(game: chess.pgn.Game) -> bool: def game_sanity_check(game: chess.pgn.Game) -> bool:
@@ -29,20 +33,51 @@ def game_sanity_check(game: chess.pgn.Game) -> bool:
return False return False
return True 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: def parse_game(game: chess.pgn.Game, writer, start_play: int=1)->None:
board: chess.Board = game.board() board: chess.Board = game.board()
if not game_sanity_check(game): if not game_sanity_check(game):
return return
result: str = game.headers["Result"] 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: if ply >= start_play:
comment: str = node.comment
writer.write("fen " + board.fen() + "\n") writer.write("fen " + board.fen() + "\n")
writer.write("move " + str(move) + "\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("ply " + str(ply)+"\n")
writer.write("result " + str(parse_result(result, board)) +"\n") writer.write("result " + str(parse_result(result, board)) +"\n")
writer.write("e\n") writer.write("e\n")
ply += 1
board.push(move) board.push(move)
def main(): def main():
@@ -53,6 +88,7 @@ def main():
args = parser.parse_args() args = parser.parse_args()
pgn_files: List[str] = glob.glob(args.pgn) 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') f = open(args.output, 'w')
for pgn_file in pgn_files: for pgn_file in pgn_files:
print("parse", pgn_file) print("parse", pgn_file)