81 lines
2.0 KiB
Python
81 lines
2.0 KiB
Python
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 |