An implementation of the Bratko-Kopec Test using python-chess

test-epd.py

#!/usr/bin/python

import chess
import chess.uci
import sys
import time


def test_epd(engine, epd):
    position = chess.Board()
    epd_info = position.set_epd(epd)

    engine.ucinewgame()
    engine.position(position)

    score = 0
    enginemove, pondermove = engine.go(movetime=60000)

    if "am" in epd_info:
        print "%s (avoid %s):" % (epd_info["id"], " ".join(map(position.san, epd_info["am"]))),
    elif "bm" in epd_info:
        print "%s (expect %s):" % (epd_info["id"], " ".join(map(position.san, epd_info["bm"]))),
    else:
        print "%s:" % epd_info["id"],

    if "am" in epd_info and enginemove in epd_info["am"]:
        pass #fail
    elif "bm" in epd_info and not enginemove in epd_info["bm"]:
        pass #fail
    else:
        score = 1
    print " %s | +%d" % (position.parse_san(enginemove), score)
    return score


def test_epd_with_fractional_scores(engine, epd):
    info_handler = chess.uci.InfoHandler()
    engine.info_handlers.append(info_handler)

    position = chess.Board()
    epd_info = position.set_epd(epd)

    engine.ucinewgame()
    engine.position(position)

    # Search in background
    engine.go(infinite=True)

    score = 0.0
    if "am" in epd_info:
        print "%s (avoid %s):" % (epd_info["id"], " ".join(map(position.san, epd_info["am"]))),
    elif "bm" in epd_info:
        print "%s (expect %s):" % (epd_info["id"], " ".join(map(position.san, epd_info["bm"]))),
    else:
        print "%s:" % epd_info["id"],

    for step in xrange(0, 3):
        time.sleep(15)

        # Assess the current principal variation.
        with info_handler as info:
            if 1 in info["pv"] and len(info["pv"][1]) >= 1:
                move = info["pv"][1][0]
                print "(%s)" % position.san(move),
                if "am" in epd_info and move in epd_info["am"]:
                    continue #fail
                elif "bm" in epd_info and not move in epd_info["bm"]:
                    continue #fail
                else:
                    score = 1.0 / (4 - step)
            else:
                print "(no pv)",

    # Assess the final best move by the engine.
    time.sleep(15)
    enginemove, pondermove = engine.stop()
    if "am" in epd_info and enginemove in epd_info["am"]:
        pass #fail
    elif "bm" in epd_info and not enginemove in epd_info["bm"]:
        pass #fail
    else:
        score = 1.0

    print "%s | +%g" % (position.san(enginemove), score)

    engine.info_handlers.remove(info_handler)
    return score


if __name__ == "__main__":
    engine = chess.uci.popen_engine(sys.argv[1])
    engine.uci()
    engine.setoption({"Threads": 2})

    score = 0.0
    count = 0
    with open(sys.argv[2],'r') as f:
        for epd in f:
            count += 1
            score += test_epd_with_fractional_scores(engine, epd)
    f.closed

    engine.quit()

    print "-------------------------------"
    print "%g / %d" % (score, count)