diff --git a/scripts/format.py b/scripts/format.py new file mode 100644 index 0000000..7846c0b --- /dev/null +++ b/scripts/format.py @@ -0,0 +1,81 @@ +import json +import os +from datetime import datetime + + +""" + A class to format a series of chess moves into specific formats for + storage. The base class expects a series of UCI moves in order. The + formatters format the UCI moves into specific output notations for + other systems to load and analyze. +""" +class BaseGameFormatter: + def __init__( + self, + save_path=None, + engine="", + white="", + black="", + result="", + strategies=None + ): + self.save_path = save_path + self.engine = engine + self.white = white + self.black = black + self.strategies = strategies + self.result = result + + self.formatted_moves = [] + self.fens = [] + + + def save_game(self, id, moves, fens): + self.formatted_moves = self.format_moves(moves) + self.fens = fens + self.save_to_disk(id) + + + def format_moves(self, moves): + """ + Define in child class. + """ + raise NotImplementedError + + + def save_to_disk(self, id): + date = datetime.today().strftime('%Y-%m-%d') + filename = f"{date}_{id}.json" + filepath = os.path.join(self.save_path, filename) + + + data = { + "date": datetime.today().strftime('%Y-%m-%d'), + "id": id, + "engine": "", + "white": "", + "black": "", + "strategies": self.strategies, + "result": self.result, + "len_moves": len(self.formatted_moves), + "moves": self.formatted_moves, + "fens": self.fens + } + + with open(filepath, "w") as out: + json.dump(data, out, indent=2) + + + def print_moves(self): + for move in self.formatted_moves: + print(move) + + +class LongPGNFormatter(BaseGameFormatter): + def format_moves(self, moves): + pgn_long = [] + it = iter(moves) + for white in it: + black = next(it, None) + pgn_long.append(f"{white} {black}" if black else white) + return pgn_long \ No newline at end of file