Analyze chess (variant) positions from an EPD file

analysis.py

import argparse
import logging
import chess, chess.variant, chess.uci


logger = logging.getLogger()
logger.setLevel(logging.INFO)


def parse_args():
    parser = argparse.ArgumentParser(description='Analyze a set of positions from an EPD file.')
    parser.add_argument('-e', '--engine', type=str, required=True, help='Path to engine.')
    parser.add_argument('-i', '--input', dest='input_file', type=str, required=True, help='Input EPD file.')
    parser.add_argument('-v', '--variant', type=str, default='chess', help='Variant name.')
    parser.add_argument('--chess960', action='store_true', help='Use chess 960 rules.')
    parser.add_argument('-p', '--print', dest='print_stdout', action='store_true', help='Print log to screen.')
    parser.add_argument('-l', '--log', help='Output log file.')
    parser.add_argument('-d', '--depth', type=int, help='Search depth.')
    parser.add_argument('-t', '--time', type=int, help='Search time in milliseconds.')
    parser.add_argument('-m', '--memory', type=int, default=16, help='Hash size in MB.')
    parser.add_argument('--threads', type=int, default=1, help='Number of threads.')
    parser.add_argument('--multipv', type=int, default=1, help='Number of multipv lines.')
    args = parser.parse_args()
    if not (args.log or args.print_stdout):
        parser.error('No logging of output defined. Add -l/--log and/or -p/--print.')
    return args


def main(args):

    if args.print_stdout:
        sh = logging.StreamHandler()
        sh.setLevel(logging.INFO)
        logger.addHandler(sh)
    if args.log is not None:
        fh = logging.FileHandler(args.log)
        fh.setLevel(logging.INFO)
        logger.addHandler(fh)

    board = chess.variant.find_variant(args.variant)()

    engine = chess.uci.popen_engine(args.engine)
    class HandlerWithLogging(chess.uci.InfoHandler):
        def pre_info(self, line):
            if not line.startswith('string'):
                logger.info(line)
            super(HandlerWithLogging, self).pre_info(line)
    info_handler = HandlerWithLogging()
    engine.info_handlers.append(info_handler)

    engine.uci()
    engine.setoption({'UCI_Variant': board.uci_variant, 'UCI_Chess960': args.chess960,
                      'Hash': args.memory, 'Threads': args.threads, 'multipv': args.multipv})

    with open(args.input_file, 'r') as epd_file:
        print('Starting to analyze positions from %s.\n' \
              'Settings: variant=%s, chess960=%s, depth=%s, movetime=%s, hash=%dMB, threads=%d, multipv=%d.'
              % (args.input_file, args.variant, args.chess960, args.depth, args.time, args.memory, args.threads, args.multipv))
        for i, epd in enumerate(epd_file):
            epd = epd.rstrip()
            print('Starting analysis %i ...' % (i + 1))
            logger.info('Analysis results of position %s\n' \
                        'for variant %s (chess960=%s) with depth=%s, movetime=%s, hash=%dMB, threads=%d, multipv=%d:'
                        % (epd, args.variant, args.chess960, args.depth, args.time, args.memory, args.threads, args.multipv))
            board.set_epd(epd)
            engine.position(board)
            engine.go(depth=args.depth, movetime=args.time)
            logger.info('')
        print('Done.')


if __name__ == "__main__":
    main(parse_args())