Stockfish self-play game generation script

stockfish_selfplay.py

import argparse

import chess
import chess.engine
import chess.pgn


def main(stockfish_path, output_path):
    engine = chess.engine.SimpleEngine.popen_uci(stockfish_path)

    game = chess.pgn.Game()
    node = game
    board = chess.Board()
    while not board.is_game_over():
        result = engine.play(board, chess.engine.Limit(time=0.5),
            info=chess.engine.Info.SCORE)
        board.push(result.move)
        node = node.add_main_variation(result.move)
        node.set_eval(result.info['score'])

    engine.quit()

    with open(output_path, 'w', encoding='utf-8') as pgn_file:
        exporter = chess.pgn.FileExporter(pgn_file, headers=False)
        pgn_string = game.accept(exporter)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--stockfish_path',
                        help='Path to Stockfish executable',
                        required=True)
    parser.add_argument('-o', '--output_path',
                        help='Path to output pgn',
                        required=True)
    args = parser.parse_args()
    main(args.stockfish_path, args.output_path)