mergepvs.py

import argparse, chess, re


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Merge possibly better PVs from several .epd files into an existing .epd file. Output is to stdout.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--preferNew",
        action="store_true",
        help="Prefer newer PVs of same length.",
    )
    parser.add_argument("source", help="The source .epd file.")
    parser.add_argument(
        "references",
        nargs="*",
        help="List of .epd files with possibly more or longer PVs.",
    )
    args = parser.parse_args()
    p = re.compile(r"^([1-8a-zA-Z/]+ [wb] [a-zA-Z\-]+ [a-h1-8\-]+)( bm #(-?\d+);)?")

    d = {}  # the dict will hold the shortest mates, with longest PVs
    for filename in [args.source] + args.references:  # important to also read in source
        with open(filename) as f:
            for line in f:
                if line.startswith("#"):  # ignore comments
                    continue
                m = p.match(line)
                assert m, f"error for line '{line[:-1]}' in file {filename}"
                fen = m.group(1)
                bm = int(m.group(3)) if m.group(2) is not None else None
                if bm is None:
                    continue
                _, _, pv = line.partition("; PV: ")
                pv, _, _ = pv[:-1].partition(";")  # remove '\n'
                pv = pv.split()
                bmold, pvold = d.get(fen, (None, None))
                if (
                    bmold is None
                    or abs(bm) < abs(bmold)
                    or bm == bmold
                    and (
                        len(pv) > len(pvold)
                        or (args.preferNew and len(pv) == len(pvold))
                    )
                ):
                    d[fen] = bm, pv

    with open(args.source) as f:
        for line in f:
            bm, pv = None, None
            if not line.startswith("#"):
                m = p.match(line)
                fen = m.group(1)
                bmold = int(m.group(3)) if m.group(2) is not None else None
                bm, pv = d.get(fen, (None, None))
            if pv:
                print(f"{fen} bm #{bm}; PV: {' '.join(pv)};")
            elif bm is not None and (bmold is None or abs(bm) < abs(bmold)):
                print(f"{fen} bm #{bm};")
            else:
                print(line[:-1])